add test for short rlp in block witness

stateless/multi_keys.nim

@@ -28,8 +28,14 @@ type
     mask*: uint
     groups*: array[16, Group]
 
-  AccountKey* = tuple[address: EthAddress, codeTouched: bool, storageKeys: MultikeysRef]
-  MatchGroup* = tuple[match: bool, group: Group]
+  AccountKey* = object
+    address*: EthAddress
+    codeTouched*: bool
+    storageKeys*: MultikeysRef
+
+  MatchGroup* = object
+    match*: bool
+    group*: Group
 
 func cmpHash(a, b: KeyHash): int =
   var i = 0
@@ -128,12 +134,12 @@ func groups*(m: MultikeysRef, depth: int, n: NibblesSeq, parentGroup: Group): Ma
       if not compareNibbles(m.keys[i].hash, depth, n):
         g.last = i - 1
         # case 1: match and no match
-        return (true, g)
+        return MatchGroup(match: true, group: g)
       inc i
 
     # case 2: all is a match group
     g.last = parentGroup.last
-    return (true, g)
+    return MatchGroup(match: true, group: g)
 
   # no match came first, skip no match
   # we only interested in a match group
@@ -149,15 +155,15 @@ func groups*(m: MultikeysRef, depth: int, n: NibblesSeq, parentGroup: Group): Ma
       if not compareNibbles(m.keys[i].hash, depth, n):
         # case 3: no match, match, and no match
         g.last = i - 1
-        return (true, g)
+        return MatchGroup(match: true,  group: g)
       inc i
 
     # case 4: no match and match
     g.last = parentGroup.last
-    return (true, g)
+    return MatchGroup(match: true, group: g)
 
   # case 5: no match at all
-  result = (false, g)
+  result = MatchGroup(match: false, group: g)
 
 func isValidMatch(mg: MatchGroup): bool {.inline.} =
   result = mg.match and mg.group.first == mg.group.last

stateless/test_witness_keys.nim

@@ -75,18 +75,18 @@ proc runTest(numPairs: int, testStatusIMPL: var TestStatus,
 
   for i in 0..<numPairs:
     let acc  = randAccount(memDB)
-    addrs[i] = (randAddress(), acc.codeTouched, acc.storageKeys)
+    addrs[i] = AccountKey(address: randAddress(), codeTouched: acc.codeTouched, storageKeys: acc.storageKeys)
     accs[i]  = acc.account
     trie.put(addrs[i].address, rlp.encode(accs[i]))
 
   when addInvalidKeys:
     # invalidAddress should not end up in block witness
     let invalidAddress = randAddress()
-    addrs.add((invalidAddress, false, MultikeysRef(nil)))
+    addrs.add(AccountKey(address: invalidAddress))
 
   if addIdenticalKeys:
     let invalidAddress = addrs[0].address
-    addrs.add((invalidAddress, false, MultikeysRef(nil)))
+    addrs.add(AccountKey(address: invalidAddress))
 
   var mkeys = newMultiKeys(addrs)
   let rootHash = trie.rootHash
@@ -122,13 +122,20 @@ proc runTest(numPairs: int, testStatusIMPL: var TestStatus,
     for kd in mkeys.keys:
       check kd.visited == true
 
-proc initMultiKeys(keys: openArray[string]): MultikeysRef =
+proc initMultiKeys(keys: openArray[string], storageMode: bool = false): MultikeysRef =
   result.new
-  for x in keys:
-    result.keys.add KeyData(
-      storageMode: false,
-      hash: hexToByteArray[32](x)
-    )
+  if storageMode:
+    for i, x in keys:
+      result.keys.add KeyData(
+        storageMode: true,
+        hash: hexToByteArray[32](x)
+      )
+  else:
+    for x in keys:
+      result.keys.add KeyData(
+        storageMode: false,
+        hash: hexToByteArray[32](x)
+      )
 
 proc parseInvalidInput(payload: openArray[byte]): bool =
   var db = newMemoryDB()
@@ -246,5 +253,34 @@ proc witnessKeysMain*() =
         let z = readFile(x)
         check parseInvalidInput(z.toOpenArrayByte(0, z.len-1))
 
+    test "short rlp test":
+      let keys = [
+        "01234567abce7762869be690036144c12c256bdb06ee9073ad5ecca18a47c325",
+        "01234567b491732f964182ce4bde5e2468318692ed446e008f621b26f8ff5660",
+        "01234567c140158288775c8912aed274fb9d6a3a260e9e95e03e70ba8df30f6b",
+      ]
+      let m  = initMultiKeys(keys, true)
+      var memDB = newMemoryDB()
+      var trie = initSecureHexaryTrie(memDB)
+      var acc  = randAccount(memDB)
+
+      var tt = initHexaryTrie(memDB)
+      for x in m.keys:
+        tt.put(x.hash, rlp.encode(1.u256))
+      acc.account.storageRoot = tt.rootHash
+
+      let addrs = @[AccountKey(address: randAddress(), codeTouched: acc.codeTouched, storageKeys: m)]
+
+      trie.put(addrs[0].address, rlp.encode(acc.account))
+      var mkeys = newMultiKeys(addrs)
+      let rootHash = trie.rootHash
+
+      var wb = initWitnessBuilder(memDB, rootHash, {wfEIP170})
+      var witness = wb.buildWitness(mkeys)
+      var db = newMemoryDB()
+      var tb = initTreeBuilder(witness, db, {wfEIP170})
+      let root = tb.buildTree()
+      check root.data == rootHash.data
+
 when isMainModule:
   witnessKeysMain()

stateless/tree_from_witness.nim

@@ -170,7 +170,8 @@ proc toNodeKey(t: var TreeBuilder, z: openArray[byte]): NodeKey =
     t.db.put(result.data, z)
 
 proc toNodeKey(z: openArray[byte]): NodeKey =
-  doAssert(z.len < 32)
+  if z.len >= 32:
+    raise newException(ParsingError, "Failed when try to convert short rlp to NodeKey")
   result.usedBytes = z.len
   result.data[0..z.len-1] = z[0..z.len-1]
 
@@ -230,7 +231,7 @@ proc buildForest*(t: var TreeBuilder): seq[KeccakHash]
     result.add KeccakHash(data: res.data)
 
 proc treeNode(t: var TreeBuilder, depth: int, storageMode = false): NodeKey =
-  if depth >= 64:
+  if depth > 64:
     raise newException(ParsingError, "invalid trie structure")
 
   let nodeType = safeReadEnum(t, TrieNodeType)
@@ -494,13 +495,13 @@ proc hashNode(t: var TreeBuilder, depth: int, storageMode: bool): NodeKey =
   if storageMode and depth >= 9:
     let z = t.safeReadByte()
     if z == ShortRlpPrefix:
-      let y = t.safeReadByte().int
-      if y == 0:
+      let rlpLen = t.safeReadByte().int
+      if rlpLen == 0:
         safeReadBytes(t, 31):
           result.toKeccak(0, t.read(31))
       else:
-        safeReadBytes(t, y):
-          result = toNodeKey(t.read(y))
+        safeReadBytes(t, rlpLen):
+          result = toNodeKey(t.read(rlpLen))
     else:
       safeReadBytes(t, 31):
         result.toKeccak(z, t.read(31))

stateless/witness_from_tree.nim

@@ -148,7 +148,7 @@ proc writeHashNode(wb: var WitnessBuilder, node: openArray[byte], depth: int, st
 
 proc writeShortRlp(wb: var WitnessBuilder, node: openArray[byte], depth: int, storageMode: bool) =
   doAssert(node.len < 32 and depth >= 9 and storageMode)
-  debugEcho "SHORT RLP ", node.len
+  wb.writeByte(HashNodeType)
   wb.writeByte(ShortRlpPrefix)
   wb.writeByte(node.len)
   wb.write(node)