deals with short node len < 32

stateless/randutils.nim

@@ -0,0 +1,59 @@
+import random, sets, nimcrypto/[utils, sysrand]
+
+type
+  RandGen*[T] = object
+    minVal, maxVal: T
+
+  Bytes* = seq[byte]
+
+proc rng*[T](minVal, maxVal: T): RandGen[T] =
+  doAssert(minVal <= maxVal)
+  result.minVal = minVal
+  result.maxVal = maxVal
+
+proc rng*[T](minMax: T): RandGen[T] =
+  rng(minMax, minMax)
+
+proc getVal*[T](x: RandGen[T]): T =
+  if x.minVal == x.maxVal: return x.minVal
+  rand(x.minVal..x.maxVal)
+
+proc randString*(len: int): string =
+  result = newString(len)
+  discard randomBytes(result[0].addr, len)
+  
+proc randBytes*(len: int): Bytes =
+  result = newSeq[byte](len)
+  discard randomBytes(result[0].addr, len)
+  
+proc randPrimitives*[T](val: int): T =
+  type
+    ByteLike = uint8 | byte | char
+
+  when T is string:
+    randString(val)
+  elif T is int:
+    result = val
+  elif T is ByteLike:
+    result = T(val)
+  elif T is Bytes:
+    result = randBytes(val)
+
+proc randList*(T: typedesc, fillGen: RandGen, listLen: int, unique: static[bool] = true): seq[T] =
+  result = newSeqOfCap[T](listLen)
+  when unique:
+    var set = initHashSet[T]()
+    for len in 0..<listLen:
+      while true:
+        let x = randPrimitives[T](fillGen.getVal())
+        if x notin set:
+          result.add x
+          set.incl x
+          break
+  else:
+    for len in 0..<listLen:
+      let x = randPrimitives[T](fillGen.getVal())
+      result.add x
+
+proc randList*(T: typedesc, fillGen, listGen: RandGen, unique: static[bool] = true): seq[T] =
+  randList(T, fillGen, listGen.getVal(), unique)

stateless/test_block_witness.nim

@@ -35,7 +35,7 @@ proc testGetBranch(tester: Tester, rootHash: KeccakHash, testStatusIMPL: var Tes
       var db = newMemoryDB()
       var tb = initTreeBuilder(witness, db)
       var root = tb.treeNode()
-      check root == rootHash
+      check root.data == rootHash.data
       #echo "ROOT: ", root.data.toHex
       #echo "rootHash: ", rootHash.data.toHex
 

stateless/test_witness_keys.nim

@@ -0,0 +1,48 @@
+import
+  randutils, stew/byteutils, random,
+  eth/[common, rlp], eth/trie/[hexary, db, trie_defs],
+  ../stateless/[witness_from_tree, tree_from_witness]
+
+proc runTest(keyBytes: int, valBytes: int, numPairs: int) =
+  var memDB = newMemoryDB()
+  var trie = initHexaryTrie(memDB)
+
+  var
+    keys = newSeq[Bytes](numPairs)
+    vals = newSeq[Bytes](numPairs)
+
+  for i in 0..<numPairs:
+    keys[i] = randList(byte, rng(0, 255), keyBytes, unique = false)
+    vals[i] = randList(byte, rng(0, 255), valBytes, unique = false)
+    trie.put(keys[i], vals[i])
+
+  let rootHash = trie.rootHash
+
+  var wb = initWitnessBuilder(memDB, rootHash)
+  var witness = wb.getBranchRecurse(keys[0])
+
+  var db = newMemoryDB()
+  var tb = initTreeBuilder(witness, db)
+  var root = tb.treeNode()
+  debugEcho "root: ", root.data.toHex
+  debugEcho "rootHash: ", rootHash.data.toHex
+  doAssert root.data == rootHash.data
+
+proc main() =
+  runTest(7, 100, 50)
+  runTest(1, 1, 1)
+  runTest(5, 5, 1)
+  runTest(6, 7, 3)
+  runTest(7, 10, 7)
+  runTest(8, 15, 11)
+  runTest(9, 30, 13)
+  runTest(11, 40, 10)
+  runTest(20, 1, 15)
+  runTest(25, 10, 20)
+  
+  randomize()
+  for i in 0..<30:
+    runTest(rand(1..30), rand(1..50), rand(1..30))
+    
+
+main()

stateless/tree_from_witness.nim

@@ -6,6 +6,10 @@ import
 type
   DB = TrieDatabaseRef
 
+  NodeKey = object
+    usedBytes: int
+    data*: array[32, byte]
+
   TreeBuilder = object
     data: seq[byte]
     pos: int
@@ -55,19 +59,35 @@ proc readU32(t: var TreeBuilder): int =
 proc toAddress(r: var EthAddress, x: openArray[byte]) {.inline.} =
   r[0..19] = x[0..19]
 
-proc toKeccak(r: var KeccakHash, x: openArray[byte]) {.inline.} =
+proc toKeccak(r: var NodeKey, x: openArray[byte]) {.inline.} =
   r.data[0..31] = x[0..31]
+  r.usedBytes = 32
 
-proc toKeccak(x: openArray[byte]): KeccakHash {.inline.} =
+proc toKeccak(x: openArray[byte]): NodeKey {.inline.} =
   result.data[0..31] = x[0..31]
+  result.usedBytes = 32
 
-proc branchNode(t: var TreeBuilder, depth: int, has16Elem: bool = true): KeccakHash
+proc append(r: var RlpWriter, n: NodeKey) =
-proc extensionNode(t: var TreeBuilder, depth: int): KeccakHash
+  if n.usedBytes < 32:
-proc accountNode(t: var TreeBuilder, depth: int): KeccakHash
+    r.append rlpFromBytes(n.data.toOpenArray(0, n.usedBytes-1))
-proc accountStorageLeafNode(t: var TreeBuilder, depth: int): KeccakHash
+  else:
-proc hashNode(t: var TreeBuilder): KeccakHash
+    r.append n.data.toOpenArray(0, n.usedBytes-1)
 
-proc treeNode*(t: var TreeBuilder, depth: int = 0, accountMode = false): KeccakHash =
+proc toNodeKey(z: openArray[byte]): NodeKey =
+  if z.len < 32:
+    result.usedBytes = z.len
+    result.data[0..z.len-1] = z[0..z.len-1]
+  else:
+    result.data = keccak(z).data
+    result.usedBytes = 32
+
+proc branchNode(t: var TreeBuilder, depth: int, has16Elem: bool = true): NodeKey
+proc extensionNode(t: var TreeBuilder, depth: int): NodeKey
+proc accountNode(t: var TreeBuilder, depth: int): NodeKey
+proc accountStorageLeafNode(t: var TreeBuilder, depth: int): NodeKey
+proc hashNode(t: var TreeBuilder): NodeKey
+
+proc treeNode*(t: var TreeBuilder, depth: int = 0, accountMode = false): NodeKey =
   assert(depth < 64)
   let nodeType = TrieNodeType(t.readByte)
 
@@ -83,9 +103,21 @@ proc treeNode*(t: var TreeBuilder, depth: int = 0, accountMode = false): KeccakH
       result = t.accountNode(depth)
   of HashNodeType: result = t.hashNode()
 
-proc branchNode(t: var TreeBuilder, depth: int, has16Elem: bool): KeccakHash =
+  if depth == 0 and result.usedBytes < 32:
+    result.data = keccak(result.data.toOpenArray(0, result.usedBytes-1)).data
+    result.usedBytes = 32
+
+proc branchNode(t: var TreeBuilder, depth: int, has16Elem: bool): NodeKey =
   assert(depth < 64)
   let mask = constructBranchMask(t.readByte, t.readByte)
+
+  when defined(debugDepth):
+    let readDepth = t.readByte.int
+    doAssert(readDepth == depth, "branchNode " & $readDepth & " vs. " & $depth)
+
+  when defined(debugHash):
+    let hash = toKeccak(t.read(32))
+
   var r = initRlpList(17)
 
   for i in 0 ..< 16:
@@ -110,7 +142,12 @@ proc branchNode(t: var TreeBuilder, depth: int, has16Elem: bool): KeccakHash =
     # anything else is empty
     r.append ""
 
-  result = keccak(r.finish)
+  result = toNodeKey(r.finish)
+
+  when defined(debugHash):
+    if result != hash:
+      debugEcho "DEPTH: ", depth
+      debugEcho "result: ", result.data.toHex, " vs. ", hash.data.toHex
 
 func hexPrefix(r: var RlpWriter, x: openArray[byte], nibblesLen: int) =
   var bytes: array[33, byte]
@@ -125,16 +162,23 @@ func hexPrefix(r: var RlpWriter, x: openArray[byte], nibblesLen: int) =
     var last = nibblesLen div 2
     for i in 1..last:
       bytes[i] = (x[i-1] shl 4) or (x[i] shr 4)
-      
+
   r.append toOpenArray(bytes, 0, nibblesLen div 2)
 
-proc extensionNode(t: var TreeBuilder, depth: int): KeccakHash =
+proc extensionNode(t: var TreeBuilder, depth: int): NodeKey =
   assert(depth < 63)
   let nibblesLen = int(t.readByte)
   assert(nibblesLen < 65)
   var r = initRlpList(2)
   r.hexPrefix(t.read(nibblesLen div 2 + nibblesLen mod 2), nibblesLen)
 
+  when defined(debugDepth):
+    let readDepth = t.readByte.int
+    doAssert(readDepth == depth, "extensionNode " & $readDepth & " vs. " & $depth)
+
+  when defined(debugHash):
+    let hash = toKeccak(t.read(32))
+
   assert(depth + nibblesLen < 65)
   let nodeType = TrieNodeType(t.readByte)
 
@@ -144,12 +188,21 @@ proc extensionNode(t: var TreeBuilder, depth: int): KeccakHash =
   of HashNodeType: r.append t.hashNode()
   else: raise newException(ValueError, "wrong type during parsing child of extension node")
 
-  result = keccak(r.finish)
+  result = toNodeKey(r.finish)
 
-proc accountNode(t: var TreeBuilder, depth: int): KeccakHash =
+  when defined(debugHash):
+    if result != hash:
+      debugEcho "DEPTH: ", depth
+    doAssert(result == hash, "EXT HASH DIFF " & result.data.toHex & " vs. " & hash.data.toHex)
+
+proc accountNode(t: var TreeBuilder, depth: int): NodeKey =
   assert(depth < 65)
   let len = t.readU32()
-  t.writeNode(t.read(len))
+  result = toNodeKey(t.read(len))
+
+  when defined(debugDepth):
+    let readDepth = t.readByte.int
+    doAssert(readDepth == depth, "accountNode " & $readDepth & " vs. " & $depth)
 
   #[let nodeType = AccountType(t.readByte)
   let nibblesLen = 64 - depth
@@ -165,12 +218,12 @@ proc accountNode(t: var TreeBuilder, depth: int): KeccakHash =
     # and reset the depth
     t.treeNode(0, accountMode = true)]#
 
-proc accountStorageLeafNode(t: var TreeBuilder, depth: int): KeccakHash =
+proc accountStorageLeafNode(t: var TreeBuilder, depth: int): NodeKey =
   assert(depth < 65)
   let nibblesLen = 64 - depth
   let pathNibbles = @(t.read(nibblesLen div 2 + nibblesLen mod 2))
   let key = @(t.read(32))
   let val = @(t.read(32))
 
-proc hashNode(t: var TreeBuilder): KeccakHash =
+proc hashNode(t: var TreeBuilder): NodeKey =
   result.toKeccak(t.read(32))

stateless/witness_from_tree.nim

@@ -29,7 +29,7 @@ proc expectHash(r: Rlp): seq[byte] =
 
 template getNode(elem: untyped): untyped =
   if elem.isList: @(elem.rawData)
-  else: get(db, elem.expectHash)
+  else: get(wb.db, elem.expectHash)
 
 proc rlpListToBitmask(r: var Rlp): uint =
   var i = 0
@@ -56,13 +56,19 @@ proc writeNibbles(wb: var WitnessBuilder; n: NibblesSeq) =
   # write nibbles
   wb.output.append(bytes.toOpenArray(0, numBytes-1))
 
-proc writeExtensionNode(wb: var WitnessBuilder, n: NibblesSeq) =
+proc writeExtensionNode(wb: var WitnessBuilder, n: NibblesSeq, depth: int, node: openArray[byte]) =
   # write type
   wb.output.append(ExtensionNodeType.byte)
   # write nibbles
   wb.writeNibbles(n)
 
-proc writeBranchNode(wb: var WitnessBuilder, mask: uint) =
+  when defined(debugDepth):
+    wb.output.append(depth.byte)
+
+  when defined(debugHash):
+    wb.output.append(keccak(node).data)
+
+proc writeBranchNode(wb: var WitnessBuilder, mask: uint, depth: int, node: openArray[byte]) =
   # write type
   if mask.branchMaskBitIsSet(16):
     wb.output.append(Branch17NodeType.byte)
@@ -72,18 +78,53 @@ proc writeBranchNode(wb: var WitnessBuilder, mask: uint) =
   wb.output.append(((mask shr 8) and 0xFF).byte)
   wb.output.append((mask and 0xFF).byte)
 
-proc writeAccountNode(wb: var WitnessBuilder, node: openArray[byte]) =
+  when defined(debugDepth):
+    wb.output.append(depth.byte)
+
+  when defined(debugHash):
+    wb.output.append(keccak(node).data)
+
+proc writeAccountNode(wb: var WitnessBuilder, node: openArray[byte], depth: int) =
   # write type
   wb.output.append(AccountNodeType.byte)
   wb.output.append(toBytesLe(node.len.uint32))
   wb.output.append(node)
 
+  when defined(debugDepth):
+    wb.output.append(depth.byte)
+
 proc writeHashNode(wb: var WitnessBuilder, node: openArray[byte]) =
   # write type
   wb.output.append(HashNodeType.byte)
   wb.output.append(node)
 
-proc getBranchRecurseAux(wb: var WitnessBuilder; db: DB, node: openArray[byte], path: NibblesSeq) =
+proc writeShortNode(wb: var WitnessBuilder, node: openArray[byte], depth: int) =
+  var nodeRlp = rlpFromBytes node
+  if not nodeRlp.hasData or nodeRlp.isEmpty: return
+  case nodeRlp.listLen
+  of 2:
+    let (isLeaf, k) = nodeRlp.extensionNodeKey
+    if isLeaf:
+      writeAccountNode(wb, node, depth)
+    else:
+      writeExtensionNode(wb, k, depth, node)
+  of 17:
+    let branchMask = rlpListToBitmask(nodeRlp)
+    writeBranchNode(wb, branchMask, depth, node)
+
+    for i in 0..<16:
+      if branchMask.branchMaskBitIsSet(i):
+        var branch = nodeRlp.listElem(i)
+        let nextLookup = branch.getNode
+        writeShortNode(wb, nextLookup, depth + 1)
+
+    var lastElem = nodeRlp.listElem(16)
+    if not lastElem.isEmpty:
+      writeAccountNode(wb, lastElem.toBytes, depth)
+  else:
+    raise newException(CorruptedTrieDatabase, "Bad Short Node")
+
+proc getBranchRecurseAux(wb: var WitnessBuilder, node: openArray[byte], path: NibblesSeq, depth: int) =
   var nodeRlp = rlpFromBytes node
   if not nodeRlp.hasData or nodeRlp.isEmpty: return
 
@@ -95,48 +136,54 @@ proc getBranchRecurseAux(wb: var WitnessBuilder; db: DB, node: openArray[byte],
       let value = nodeRlp.listElem(1)
       if not isLeaf:
         # ExtensionNodeType
-        writeExtensionNode(wb, k)
+        writeExtensionNode(wb, k, depth, node)
         let nextLookup = value.getNode
-        getBranchRecurseAux(wb, db, nextLookup, path.slice(sharedNibbles))
+        getBranchRecurseAux(wb, nextLookup, path.slice(sharedNibbles), depth + sharedNibbles)
       else:
         # AccountNodeType
-        writeAccountNode(wb, value.toBytes)
+        writeAccountNode(wb, node, depth)
     else:
       writeHashNode(wb, keccak(node).data)
   of 17:
     let branchMask = rlpListToBitmask(nodeRlp)
-    writeBranchNode(wb, branchMask)
+    writeBranchNode(wb, branchMask, depth, node)
+
     let notLeaf = path.len != 0
     for i in 0..<16:
       if branchMask.branchMaskBitIsSet(i):
         var branch = nodeRlp.listElem(i)
         if notLeaf and i == path[0].int:
           let nextLookup = branch.getNode
-          getBranchRecurseAux(wb, db, nextLookup, path.slice(1))
+          getBranchRecurseAux(wb, nextLookup, path.slice(1), depth + 1)
         else:
-          writeHashNode(wb, branch.expectHash)
+          if branch.isList:
+            let nextLookup = branch.getNode
+            writeShortNode(wb, nextLookup, depth + 1)
+          else:
+            writeHashNode(wb, branch.expectHash)
 
     # put 17th elem
     var lastElem = nodeRlp.listElem(16)
     if not lastElem.isEmpty:
       if path.len == 0:
-        writeAccountNode(wb, lastElem.toBytes)
+        doAssert(false, "ACC NODE A?")
+        writeAccountNode(wb, lastElem.toBytes, depth)
       else:
-        writeHashNode(wb, lastElem.toBytes)
+        doAssert(false, "HASH NODE B?")
+        writeHashNode(wb, lastElem.expectHash)
   else:
     raise newException(CorruptedTrieDatabase,
                        "HexaryTrie node with an unexpected number of children")
 
 proc getBranchRecurse*(wb: var WitnessBuilder; key: openArray[byte]): seq[byte] =
   var node = wb.db.get(wb.root.data)
-  getBranchRecurseAux(wb, wb.db, node, initNibbleRange(key))
+  getBranchRecurseAux(wb, node, initNibbleRange(key), 0)
   shallowCopy(result, wb.output.getOutput(seq[byte]))
 
-proc getBranchStack*(self: WitnessBuilder; key: openArray[byte]): string =
+proc getBranchStack*(wb: WitnessBuilder; key: openArray[byte]): string =
   var
-    node = self.db.get(self.root.data)
+    node = wb.db.get(wb.root.data)
     stack = @[(node, initNibbleRange(key))]
-    db = self.db
 
   result.add node.toHex
   while stack.len > 0:
@@ -165,27 +212,10 @@ proc getBranchStack*(self: WitnessBuilder; key: openArray[byte]): string =
       raise newException(CorruptedTrieDatabase,
                        "HexaryTrie node with an unexpected number of children")
 
-#[
+proc getBranch*(wb: WitnessBuilder; key: openArray[byte]): seq[seq[byte]] =
-proc hexPrefixDecode*(r: openArray[byte]): tuple[isLeaf: bool, nibbles: NibblesSeq] =
-  result.nibbles = initNibbleRange(r)
-  if r.len > 0:
-    result.isLeaf = (r[0] and 0x20) != 0
-    let hasOddLen = (r[0] and 0x10) != 0
-    result.nibbles.ibegin = 2 - int(hasOddLen)
-  else:
-    result.isLeaf = false
-
-proc sharedPrefixLen*(lhs, rhs: NibblesSeq): int =
-  result = 0
-  while result < lhs.len and result < rhs.len:
-    if lhs[result] != rhs[result]: break
-    inc result
-]#
-proc getBranch*(self: WitnessBuilder; key: openArray[byte]): seq[seq[byte]] =
   var
-    node = self.db.get(self.root.data)
+    node = wb.db.get(wb.root.data)
     stack = @[(node, initNibbleRange(key))]
-    db = self.db
 
   result.add node
   while stack.len > 0:
@@ -214,11 +244,10 @@ proc getBranch*(self: WitnessBuilder; key: openArray[byte]): seq[seq[byte]] =
       raise newException(CorruptedTrieDatabase,
                        "HexaryTrie node with an unexpected number of children")
 
-proc buildWitness*(self: var WitnessBuilder; key: openArray[byte]) =
+proc buildWitness*(wb: var WitnessBuilder; key: openArray[byte]) =
   var
-    node = self.db.get(self.root.data)
+    node = wb.db.get(wb.root.data)
     stack = @[(node, initNibbleRange(key))]
-    db = self.db
 
   while stack.len > 0:
     let (node, path) = stack.pop()
@@ -243,4 +272,3 @@ proc buildWitness*(self: var WitnessBuilder; key: openArray[byte]) =
     else:
       raise newException(CorruptedTrieDatabase,
                        "HexaryTrie node with an unexpected number of children")
-