thecoderpanda
diff --git a/‎trie/committer.go
Lines changed: 18 additions & 22 deletions b/‎trie/committer.go
Lines changed: 18 additions & 22 deletions
@@ -142,12 +142,10 @@ func (c *committer) store(path []byte, n node) node {
 	// We have the hash already, estimate the RLP encoding-size of the node.
 	// The size is used for mem tracking, does not need to be exact
 	var (
-		size  = estimateSize(n)
 		nhash = common.BytesToHash(hash)
 		mnode = &memoryNode{
 			hash: nhash,
-			node: simplifyNode(n),
-			size: uint16(size),
+			node: nodeToBytes(n),
 		}
 	)
 	// Collect the dirty node to nodeset for return.
@@ -166,31 +164,29 @@ func (c *committer) store(path []byte, n node) node {
 	return hash
 }
 
-// estimateSize estimates the size of an rlp-encoded node, without actually
-// rlp-encoding it (zero allocs). This method has been experimentally tried, and with a trie
-// with 1000 leaves, the only errors above 1% are on small shortnodes, where this
-// method overestimates by 2 or 3 bytes (e.g. 37 instead of 35)
-func estimateSize(n node) int {
+// mptResolver the children resolver in merkle-patricia-tree.
+type mptResolver struct{}
+
+// ForEach implements childResolver, decodes the provided node and
+// traverses the children inside.
+func (resolver mptResolver) forEach(node []byte, onChild func(common.Hash)) {
+	forGatherChildren(mustDecodeNodeUnsafe(nil, node), onChild)
+}
+
+// forGatherChildren traverses the node hierarchy and invokes the callback
+// for all the hashnode children.
+func forGatherChildren(n node, onChild func(hash common.Hash)) {
 	switch n := n.(type) {
 	case *shortNode:
-		// A short node contains a compacted key, and a value.
-		return 3 + len(n.Key) + estimateSize(n.Val)
+		forGatherChildren(n.Val, onChild)
 	case *fullNode:
-		// A full node contains up to 16 hashes (some nils), and a key
-		s := 3
 		for i := 0; i < 16; i++ {
-			if child := n.Children[i]; child != nil {
-				s += estimateSize(child)
-			} else {
-				s++
-			}
+			forGatherChildren(n.Children[i], onChild)
 		}
-		return s
-	case valueNode:
-		return 1 + len(n)
 	case hashNode:
-		return 1 + len(n)
+		onChild(common.BytesToHash(n))
+	case valueNode, nil:
 	default:
-		panic(fmt.Sprintf("node type %T", n))
+		panic(fmt.Sprintf("unknown node type: %T", n))
 	}
 }