From 82269e4c91f7d255cb66cba746abd4d634fa286b Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Wed, 13 Nov 2024 08:59:56 +0100
Subject: [PATCH] trie, core: use unsafe pool internally in stacktrie

---
 core/types/hashing_test.go |  4 +++-
 core/types/tx_blob_test.go | 18 +++++++++++++++---
 trie/bytepool.go           | 37 +++++++++++++++++++++++++++++++++++++
 trie/stacktrie.go          |  5 +++--
 4 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/core/types/hashing_test.go b/core/types/hashing_test.go
index d51b60e76c..905d157756 100644
--- a/core/types/hashing_test.go
+++ b/core/types/hashing_test.go
@@ -94,11 +94,13 @@ func BenchmarkDeriveSha200(b *testing.B) {
 		}
 	})
 
+	st := trie.NewStackTrie(nil)
 	b.Run("stack_trie", func(b *testing.B) {
 		b.ResetTimer()
 		b.ReportAllocs()
 		for i := 0; i < b.N; i++ {
-			have = types.DeriveSha(txs, trie.NewStackTrie(nil))
+			st.Reset()
+			have = types.DeriveSha(txs, st)
 		}
 		if have != want {
 			b.Errorf("have %x want %x", have, want)
diff --git a/core/types/tx_blob_test.go b/core/types/tx_blob_test.go
index 6bd0f183b7..c9def20cf0 100644
--- a/core/types/tx_blob_test.go
+++ b/core/types/tx_blob_test.go
@@ -2,6 +2,7 @@ package types
 
 import (
 	"crypto/ecdsa"
+	"sync"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -58,12 +59,22 @@ func TestBlobTxSize(t *testing.T) {
 	}
 }
 
+// emptyInit ensures that we init the kzg empties only once
 var (
-	emptyBlob          = new(kzg4844.Blob)
-	emptyBlobCommit, _ = kzg4844.BlobToCommitment(emptyBlob)
-	emptyBlobProof, _  = kzg4844.ComputeBlobProof(emptyBlob, emptyBlobCommit)
+	emptyInit       sync.Once
+	emptyBlob       *kzg4844.Blob
+	emptyBlobCommit kzg4844.Commitment
+	emptyBlobProof  kzg4844.Proof
 )
 
+func initEmpties() {
+	emptyInit.Do(func() {
+		emptyBlob = new(kzg4844.Blob)
+		emptyBlobCommit, _ = kzg4844.BlobToCommitment(emptyBlob)
+		emptyBlobProof, _ = kzg4844.ComputeBlobProof(emptyBlob, emptyBlobCommit)
+	})
+}
+
 func createEmptyBlobTx(key *ecdsa.PrivateKey, withSidecar bool) *Transaction {
 	blobtx := createEmptyBlobTxInner(withSidecar)
 	signer := NewCancunSigner(blobtx.ChainID.ToBig())
@@ -71,6 +82,7 @@ func createEmptyBlobTx(key *ecdsa.PrivateKey, withSidecar bool) *Transaction {
 }
 
 func createEmptyBlobTxInner(withSidecar bool) *BlobTx {
+	initEmpties()
 	sidecar := &BlobTxSidecar{
 		Blobs:       []kzg4844.Blob{*emptyBlob},
 		Commitments: []kzg4844.Commitment{emptyBlobCommit},
diff --git a/trie/bytepool.go b/trie/bytepool.go
index 00b567ae1d..2138c08c69 100644
--- a/trie/bytepool.go
+++ b/trie/bytepool.go
@@ -53,3 +53,40 @@ func (bp *bytesPool) Put(b []byte) {
 	default:
 	}
 }
+
+// unsafeBytesPool is a pool for byteslices. It is not safe for concurrent use.
+type unsafeBytesPool struct {
+	items [][]byte
+	w     int
+}
+
+// newUnsafeBytesPool creates a new bytesPool. The sliceCap sets the capacity of
+// newly allocated slices, and the nitems determines how many items the pool
+// will hold, at maximum.
+func newUnsafeBytesPool(sliceCap, nitems int) *unsafeBytesPool {
+	return &unsafeBytesPool{
+		items: make([][]byte, 0, nitems),
+		w:     sliceCap,
+	}
+}
+
+// Get returns a slice.
+func (bp *unsafeBytesPool) Get() []byte {
+	if len(bp.items) > 0 {
+		last := bp.items[len(bp.items)-1]
+		bp.items = bp.items[:len(bp.items)-1]
+		return last
+	}
+	return make([]byte, 0, bp.w)
+}
+
+// Put returns a slice to the pool. This method
+// will ignore slices that are too small or too large (>3x the cap)
+func (bp *unsafeBytesPool) Put(b []byte) {
+	if c := cap(b); c < bp.w || c > 3*bp.w {
+		return
+	}
+	if len(bp.items) < cap(bp.items) {
+		bp.items = append(bp.items, b)
+	}
+}
diff --git a/trie/stacktrie.go b/trie/stacktrie.go
index 05bc05e153..796f98aa69 100644
--- a/trie/stacktrie.go
+++ b/trie/stacktrie.go
@@ -50,7 +50,7 @@ type StackTrie struct {
 	onTrieNode OnTrieNode
 	kBuf       []byte // buf space used for hex-key during insertions
 	pBuf       []byte // buf space used for path during insertions
-	vPool      *bytesPool
+	vPool      *unsafeBytesPool
 }
 
 // NewStackTrie allocates and initializes an empty trie. The committed nodes
@@ -62,7 +62,7 @@ func NewStackTrie(onTrieNode OnTrieNode) *StackTrie {
 		onTrieNode: onTrieNode,
 		kBuf:       make([]byte, 0, 64),
 		pBuf:       make([]byte, 0, 32),
-		vPool:      newBytesPool(300, 20),
+		vPool:      newUnsafeBytesPool(300, 20),
 	}
 }
 
@@ -113,6 +113,7 @@ func (t *StackTrie) Update(key, value []byte) error {
 func (t *StackTrie) Reset() {
 	t.root = stPool.Get().(*stNode)
 	t.last = nil
+	t.onTrieNode = nil
 }
 
 // TrieKey returns the internal key representation for the given user key.