From 22f86a7668605416da091db2755e8fcb819b6f15 Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Mon, 11 Nov 2024 15:18:42 +0100
Subject: [PATCH] trie: stacktrie pool hashing slices

---
 trie/hasher.go    |  8 ++++++++
 trie/stacktrie.go | 28 ++++++++++++++++++++++++----
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/trie/hasher.go b/trie/hasher.go
index abf654c709..a327e3bf0a 100644
--- a/trie/hasher.go
+++ b/trie/hasher.go
@@ -188,6 +188,14 @@ func (h *hasher) hashData(data []byte) hashNode {
 	return n
 }
 
+// hashDataTo hashes the provided data to the dest buffer (must be at least
+// 32 byte large)
+func (h *hasher) hashDataTo(data []byte, dest []byte) {
+	h.sha.Reset()
+	h.sha.Write(data)
+	h.sha.Read(dest)
+}
+
 // proofHash is used to construct trie proofs, and returns the 'collapsed'
 // node (for later RLP encoding) as well as the hashed node -- unless the
 // node is smaller than 32 bytes, in which case it will be returned as is.
diff --git a/trie/stacktrie.go b/trie/stacktrie.go
index 247ab777df..b9a360e675 100644
--- a/trie/stacktrie.go
+++ b/trie/stacktrie.go
@@ -27,6 +27,7 @@ import (
 
 var (
 	stPool = sync.Pool{New: func() any { return new(stNode) }}
+	bPool  = sync.Pool{New: func() any { return make([]byte, 0, 32) }}
 	_      = types.TrieHasher((*StackTrie)(nil))
 )
 
@@ -145,6 +146,12 @@ const (
 )
 
 func (n *stNode) reset() *stNode {
+	if n.typ == hashedNode {
+		// On hashnodes, we 'own' the val: it is guaranteed to be not held
+		// by external caller. Hence, when we arrive here, we can put it back
+		// into the pool
+		bPool.Put(n.val)
+	}
 	n.key = n.key[:0]
 	n.val = nil
 	for i := range n.children {
@@ -342,11 +349,16 @@ func (t *StackTrie) hash(st *stNode, path []byte) {
 			}
 			t.hash(child, append(path, byte(i)))
 			nodes.Children[i] = child.val
-			st.children[i] = nil
-			stPool.Put(child.reset()) // Release child back to pool.
 		}
 		nodes.encode(t.h.encbuf)
 		blob = t.h.encodedBytes()
+		for i, child := range st.children {
+			if child == nil {
+				continue
+			}
+			st.children[i] = nil
+			stPool.Put(child.reset()) // Release child back to pool.
+		}
 
 	case extNode:
 		// recursively hash and commit child as the first step
@@ -381,15 +393,23 @@ func (t *StackTrie) hash(st *stNode, path []byte) {
 	st.typ = hashedNode
 	st.key = st.key[:0]
 
+	st.val = nil // Release reference to potentially externally held slice.
+
 	// Skip committing the non-root node if the size is smaller than 32 bytes
 	// as tiny nodes are always embedded in their parent except root node.
 	if len(blob) < 32 && len(path) > 0 {
-		st.val = common.CopyBytes(blob)
+		val := bPool.Get().([]byte)
+		val = val[:len(blob)]
+		copy(val, blob)
+		st.val = val
 		return
 	}
 	// Write the hash to the 'val'. We allocate a new val here to not mutate
 	// input values.
-	st.val = t.h.hashData(blob)
+	val := bPool.Get().([]byte)
+	val = val[:32]
+	t.h.hashDataTo(blob, val)
+	st.val = val
 
 	// Invoke the callback it's provided. Notably, the path and blob slices are
 	// volatile, please deep-copy the slices in callback if the contents need