diff --git a/eth/trie/hexary.nim b/eth/trie/hexary.nim index a0f88bd..2f597c3 100644 --- a/eth/trie/hexary.nim +++ b/eth/trie/hexary.nim @@ -270,6 +270,50 @@ iterator pairs*(self: HexaryTrie): (BytesRange, BytesRange) = let res = getPairsAux(self.db, stack) yield res +iterator replicate*(self: HexaryTrie): (BytesRange, BytesRange) = + # this iterator helps 'rebuild' the entire trie without + # going through a trie algorithm, but it will pull the entire + # low level KV pairs. Thus the target db will only use put operations + # without del or contains, can speed up huge trie replication. + var + localBytes = keyToLocalBytes(self.db, self.root) + nodeRlp = rlpFromBytes localBytes + path = newRange[byte](0) + stack = @[(nodeRlp, initNibbleRange(path))] + + template pushOrYield(elem: untyped) = + if elem.isList: + stack.add((elem, key)) + else: + let rlpBytes = get(self.db, toOpenArray(elem.expectHash)).toRange + let nextLookup = rlpFromBytes(rlpBytes) + stack.add((nextLookup, key)) + yield (elem.toBytes, rlpBytes) + + yield (self.rootHash.toRange, localBytes) + while stack.len > 0: + let (nodeRlp, path) = stack.pop() + if not nodeRlp.hasData or nodeRlp.isEmpty: + continue + + case nodeRlp.listLen + of 2: + let + (isLeaf, k) = nodeRlp.extensionNodeKey + key = path & k + value = nodeRlp.listElem(1) + if not isLeaf: pushOrYield(value) + of 17: + for i in 0 ..< 16: + var branch = nodeRlp.listElem(i) + if not branch.isEmpty: + var key = path.cloneAndReserveNibble() + key.replaceLastNibble(i.byte) + pushOrYield(branch) + else: + raise newException(CorruptedTrieError, + "HexaryTrie node with an unexpected number of children") + proc getValues*(self: HexaryTrie): seq[BytesRange] = result = @[] for v in self.values: diff --git a/tests/trie/test_hexary_trie.nim b/tests/trie/test_hexary_trie.nim index 0fbeae0..3ebe583 100644 --- a/tests/trie/test_hexary_trie.nim +++ b/tests/trie/test_hexary_trie.nim @@ -413,3 +413,31 @@ suite "hexary trie": for x in 0 ..< numKeyVal: var branch = trie2.getBranch(keys[x]) check isValidBranch(branch, trie2.rootHash, keys[x], newVals[x]) + + test "replicate iterator": + const + numKeyVal = 30 + + var + memdb = newMemoryDB() + repdb = newMemoryDB() + pruningTrie = initHexaryTrie(memdb, isPruning = true) + + let + keys = randList(BytesRange, randGen(5, 77), randGen(numKeyVal)) + vals = randList(BytesRange, randGen(1, 57), randGen(numKeyVal)) + + for i in 0 ..< numKeyVal: + pruningTrie.put(keys[i], vals[i]) + + let rootHash = pruningTrie.rootHash + for k, v in pruningTrie.replicate: + repdb.put(k.toOpenArray, v.toOpenArray) + + var trie = initHexaryTrie(repdb, rootHash, isPruning = true) + var numPairs = 0 + for k, v in trie.pairs: + check k in keys + check v in vals + inc numPairs + check numPairs == numKeyVal