trie: optimize memory allocation (#30932)

This pull request removes the node copy operation to reduce memory
allocation. Key Changes as below:

**(a) Use `decodeNodeUnsafe` for decoding nodes retrieved from the trie
node reader**

In the current implementation of the MPT, once a trie node blob is
retrieved, it is passed to `decodeNode` for decoding. However,
`decodeNode` assumes the supplied byte slice might be mutated later, so
it performs a deep copy internally before parsing the node.

Given that the node reader is implemented by the path database and the
hash database, both of which guarantee the immutability of the returned
byte slice. By restricting the node reader interface to explicitly
guarantee that the returned byte slice will not be modified, we can
safely replace `decodeNode` with `decodeNodeUnsafe`. This eliminates the
need for a redundant byte copy during each node resolution.

**(b) Modify the trie in place**

In the current implementation of the MPT, a copy of a trie node is
created before any modifications are made. These modifications include:
- Node resolution: Converting the value from a hash to the actual node.
- Node hashing: Tagging the hash into its cache.
- Node commit: Replacing the children with its hash.
- Structural changes: For example, adding a new child to a fullNode or
replacing a child of a shortNode.

This mechanism ensures that modifications only affect the live tree,
leaving all previously created copies unaffected.

Unfortunately, this property leads to a huge memory allocation
requirement. For example, if we want to modify the fullNode for n times,
the node will be copied for n times.

In this pull request, all the trie modifications are made in place. In
order to make sure all previously created copies are unaffected, the
`Copy` function now will deep-copy all the live nodes rather than the
root node itself.

With this change, while the `Copy` function becomes more expensive, it's
totally acceptable as it's not a frequently used one. For the normal
trie operations (Get, GetNode, Hash, Commit, Insert, Delete), the node
copy is not required anymore.

Author: rjl493456442

trie/committer.go

@@ -57,32 +57,26 @@ func (c *committer) commit(path []byte, n node, parallel bool) node {
 	// Commit children, then parent, and remove the dirty flag.
 	switch cn := n.(type) {
 	case *shortNode:
-		// Commit child
-		collapsed := cn.copy()
-
 		// If the child is fullNode, recursively commit,
 		// otherwise it can only be hashNode or valueNode.
 		if _, ok := cn.Val.(*fullNode); ok {
-			collapsed.Val = c.commit(append(path, cn.Key...), cn.Val, false)
+			cn.Val = c.commit(append(path, cn.Key...), cn.Val, false)
 		}
 		// The key needs to be copied, since we're adding it to the
 		// modified nodeset.
-		collapsed.Key = hexToCompact(cn.Key)
-		hashedNode := c.store(path, collapsed)
+		cn.Key = hexToCompact(cn.Key)
+		hashedNode := c.store(path, cn)
 		if hn, ok := hashedNode.(hashNode); ok {
 			return hn
 		}
-		return collapsed
+		return cn
 	case *fullNode:
-		hashedKids := c.commitChildren(path, cn, parallel)
-		collapsed := cn.copy()
-		collapsed.Children = hashedKids
-
-		hashedNode := c.store(path, collapsed)
+		c.commitChildren(path, cn, parallel)
+		hashedNode := c.store(path, cn)
 		if hn, ok := hashedNode.(hashNode); ok {
 			return hn
 		}
-		return collapsed
+		return cn
 	case hashNode:
 		return cn
 	default:
@@ -92,11 +86,10 @@ func (c *committer) commit(path []byte, n node, parallel bool) node {
 }
 
 // commitChildren commits the children of the given fullnode
-func (c *committer) commitChildren(path []byte, n *fullNode, parallel bool) [17]node {
+func (c *committer) commitChildren(path []byte, n *fullNode, parallel bool) {
 	var (
-		wg       sync.WaitGroup
-		nodesMu  sync.Mutex
-		children [17]node
+		wg      sync.WaitGroup
+		nodesMu sync.Mutex
 	)
 	for i := 0; i < 16; i++ {
 		child := n.Children[i]
@@ -106,22 +99,21 @@ func (c *committer) commitChildren(path []byte, n *fullNode, parallel bool) [17]
 		// If it's the hashed child, save the hash value directly.
 		// Note: it's impossible that the child in range [0, 15]
 		// is a valueNode.
-		if hn, ok := child.(hashNode); ok {
-			children[i] = hn
+		if _, ok := child.(hashNode); ok {
 			continue
 		}
 		// Commit the child recursively and store the "hashed" value.
 		// Note the returned node can be some embedded nodes, so it's
 		// possible the type is not hashNode.
 		if !parallel {
-			children[i] = c.commit(append(path, byte(i)), child, false)
+			n.Children[i] = c.commit(append(path, byte(i)), child, false)
 		} else {
 			wg.Add(1)
 			go func(index int) {
 				p := append(path, byte(index))
 				childSet := trienode.NewNodeSet(c.nodes.Owner)
 				childCommitter := newCommitter(childSet, c.tracer, c.collectLeaf)
-				children[index] = childCommitter.commit(p, child, false)
+				n.Children[index] = childCommitter.commit(p, child, false)
 				nodesMu.Lock()
 				c.nodes.MergeSet(childSet)
 				nodesMu.Unlock()
@@ -132,11 +124,6 @@ func (c *committer) commitChildren(path []byte, n *fullNode, parallel bool) [17]
 	if parallel {
 		wg.Wait()
 	}
-	// For the 17th child, it's possible the type is valuenode.
-	if n.Children[16] != nil {
-		children[16] = n.Children[16]
-	}
-	return children
 }
 
 // store hashes the node n and adds it to the modified nodeset. If leaf collection

trie/hasher.go

@@ -53,72 +53,66 @@ func returnHasherToPool(h *hasher) {
 	hasherPool.Put(h)
 }
 
-// hash collapses a node down into a hash node, also returning a copy of the
-// original node initialized with the computed hash to replace the original one.
-func (h *hasher) hash(n node, force bool) (hashed node, cached node) {
+// hash collapses a node down into a hash node.
+func (h *hasher) hash(n node, force bool) node {
 	// Return the cached hash if it's available
 	if hash, _ := n.cache(); hash != nil {
-		return hash, n
+		return hash
 	}
 	// Trie not processed yet, walk the children
 	switch n := n.(type) {
 	case *shortNode:
-		collapsed, cached := h.hashShortNodeChildren(n)
+		collapsed := h.hashShortNodeChildren(n)
 		hashed := h.shortnodeToHash(collapsed, force)
-		// We need to retain the possibly _not_ hashed node, in case it was too
-		// small to be hashed
 		if hn, ok := hashed.(hashNode); ok {
-			cached.flags.hash = hn
+			n.flags.hash = hn
 		} else {
-			cached.flags.hash = nil
+			n.flags.hash = nil
 		}
-		return hashed, cached
+		return hashed
 	case *fullNode:
-		collapsed, cached := h.hashFullNodeChildren(n)
-		hashed = h.fullnodeToHash(collapsed, force)
+		collapsed := h.hashFullNodeChildren(n)
+		hashed := h.fullnodeToHash(collapsed, force)
 		if hn, ok := hashed.(hashNode); ok {
-			cached.flags.hash = hn
+			n.flags.hash = hn
 		} else {
-			cached.flags.hash = nil
+			n.flags.hash = nil
 		}
-		return hashed, cached
+		return hashed
 	default:
 		// Value and hash nodes don't have children, so they're left as were
-		return n, n
+		return n
 	}
 }
 
-// hashShortNodeChildren collapses the short node. The returned collapsed node
-// holds a live reference to the Key, and must not be modified.
-func (h *hasher) hashShortNodeChildren(n *shortNode) (collapsed, cached *shortNode) {
-	// Hash the short node's child, caching the newly hashed subtree
-	collapsed, cached = n.copy(), n.copy()
-	// Previously, we did copy this one. We don't seem to need to actually
-	// do that, since we don't overwrite/reuse keys
-	// cached.Key = common.CopyBytes(n.Key)
+// hashShortNodeChildren returns a copy of the supplied shortNode, with its child
+// being replaced by either the hash or an embedded node if the child is small.
+func (h *hasher) hashShortNodeChildren(n *shortNode) *shortNode {
+	var collapsed shortNode
 	collapsed.Key = hexToCompact(n.Key)
-	// Unless the child is a valuenode or hashnode, hash it
 	switch n.Val.(type) {
 	case *fullNode, *shortNode:
-		collapsed.Val, cached.Val = h.hash(n.Val, false)
+		collapsed.Val = h.hash(n.Val, false)
+	default:
+		collapsed.Val = n.Val
 	}
-	return collapsed, cached
+	return &collapsed
 }
 
-func (h *hasher) hashFullNodeChildren(n *fullNode) (collapsed *fullNode, cached *fullNode) {
-	// Hash the full node's children, caching the newly hashed subtrees
-	cached = n.copy()
-	collapsed = n.copy()
+// hashFullNodeChildren returns a copy of the supplied fullNode, with its child
+// being replaced by either the hash or an embedded node if the child is small.
+func (h *hasher) hashFullNodeChildren(n *fullNode) *fullNode {
+	var children [17]node
 	if h.parallel {
 		var wg sync.WaitGroup
 		wg.Add(16)
 		for i := 0; i < 16; i++ {
 			go func(i int) {
 				hasher := newHasher(false)
 				if child := n.Children[i]; child != nil {
-					collapsed.Children[i], cached.Children[i] = hasher.hash(child, false)
+					children[i] = hasher.hash(child, false)
 				} else {
-					collapsed.Children[i] = nilValueNode
+					children[i] = nilValueNode
 				}
 				returnHasherToPool(hasher)
 				wg.Done()
@@ -128,19 +122,21 @@ func (h *hasher) hashFullNodeChildren(n *fullNode) (collapsed *fullNode, cached
 	} else {
 		for i := 0; i < 16; i++ {
 			if child := n.Children[i]; child != nil {
-				collapsed.Children[i], cached.Children[i] = h.hash(child, false)
+				children[i] = h.hash(child, false)
 			} else {
-				collapsed.Children[i] = nilValueNode
+				children[i] = nilValueNode
 			}
 		}
 	}
-	return collapsed, cached
+	if n.Children[16] != nil {
+		children[16] = n.Children[16]
+	}
+	return &fullNode{flags: nodeFlag{}, Children: children}
 }
 
-// shortnodeToHash creates a hashNode from a shortNode. The supplied shortnode
-// should have hex-type Key, which will be converted (without modification)
-// into compact form for RLP encoding.
-// If the rlp data is smaller than 32 bytes, `nil` is returned.
+// shortNodeToHash computes the hash of the given shortNode. The shortNode must
+// first be collapsed, with its key converted to compact form. If the RLP-encoded
+// node data is smaller than 32 bytes, the node itself is returned.
 func (h *hasher) shortnodeToHash(n *shortNode, force bool) node {
 	n.encode(h.encbuf)
 	enc := h.encodedBytes()
@@ -151,8 +147,8 @@ func (h *hasher) shortnodeToHash(n *shortNode, force bool) node {
 	return h.hashData(enc)
 }
 
-// fullnodeToHash is used to create a hashNode from a fullNode, (which
-// may contain nil values)
+// fullnodeToHash computes the hash of the given fullNode. If the RLP-encoded
+// node data is smaller than 32 bytes, the node itself is returned.
 func (h *hasher) fullnodeToHash(n *fullNode, force bool) node {
 	n.encode(h.encbuf)
 	enc := h.encodedBytes()
@@ -203,10 +199,10 @@ func (h *hasher) hashDataTo(dst, data []byte) {
 func (h *hasher) proofHash(original node) (collapsed, hashed node) {
 	switch n := original.(type) {
 	case *shortNode:
-		sn, _ := h.hashShortNodeChildren(n)
+		sn := h.hashShortNodeChildren(n)
 		return sn, h.shortnodeToHash(sn, false)
 	case *fullNode:
-		fn, _ := h.hashFullNodeChildren(n)
+		fn := h.hashFullNodeChildren(n)
 		return fn, h.fullnodeToHash(fn, false)
 	default:
 		// Value and hash nodes don't have children, so they're left as were

trie/node.go

@@ -79,15 +79,19 @@ func (n *fullNode) EncodeRLP(w io.Writer) error {
 	return eb.Flush()
 }
 
-func (n *fullNode) copy() *fullNode   { copy := *n; return &copy }
-func (n *shortNode) copy() *shortNode { copy := *n; return &copy }
-
 // nodeFlag contains caching-related metadata about a node.
 type nodeFlag struct {
 	hash  hashNode // cached hash of the node (may be nil)
 	dirty bool     // whether the node has changes that must be written to the database
 }
 
+func (n nodeFlag) copy() nodeFlag {
+	return nodeFlag{
+		hash:  common.CopyBytes(n.hash),
+		dirty: n.dirty,
+	}
+}
+
 func (n *fullNode) cache() (hashNode, bool)  { return n.flags.hash, n.flags.dirty }
 func (n *shortNode) cache() (hashNode, bool) { return n.flags.hash, n.flags.dirty }
 func (n hashNode) cache() (hashNode, bool)   { return nil, true }
@@ -228,7 +232,9 @@ func decodeRef(buf []byte) (node, []byte, error) {
 			err := fmt.Errorf("oversized embedded node (size is %d bytes, want size < %d)", size, hashLen)
 			return nil, buf, err
 		}
-		n, err := decodeNode(nil, buf)
+		// The buffer content has already been copied or is safe to use;
+		// no additional copy is required.
+		n, err := decodeNodeUnsafe(nil, buf)
 		return n, rest, err
 	case kind == rlp.String && len(val) == 0:
 		// empty node