From 1b00b6e51d12a7a79e6fdb8bde3d0b0a8765b050 Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Mon, 11 Nov 2024 22:54:26 +0100
Subject: [PATCH] trie: implement bytepool

---
 trie/bytepool.go  | 39 +++++++++++++++++++++++++++++++++++++++
 trie/stacktrie.go |  6 +++---
 2 files changed, 42 insertions(+), 3 deletions(-)
 create mode 100644 trie/bytepool.go

diff --git a/trie/bytepool.go b/trie/bytepool.go
new file mode 100644
index 0000000000..8f331a49d5
--- /dev/null
+++ b/trie/bytepool.go
@@ -0,0 +1,39 @@
+package trie
+
+type bytepool struct {
+	c chan []byte
+	w int
+	h int
+}
+
+func newByteslicepool(sliceCap, nitems int) *bytepool {
+	b := &bytepool{
+		c: make(chan []byte, nitems),
+		w: sliceCap,
+	}
+	return b
+}
+
+func (bp *bytepool) Get() []byte {
+	select {
+	case b := <-bp.c:
+		return b
+	default:
+		return make([]byte, 0, bp.w)
+	}
+}
+
+func (bp *bytepool) Put(b []byte) {
+	// Ignore too small slices
+	if cap(b) < bp.w {
+		return
+	}
+	// Don't retain too large slices either
+	if cap(b) > 3*bp.w {
+		return
+	}
+	select {
+	case bp.c <- b:
+	default:
+	}
+}
diff --git a/trie/stacktrie.go b/trie/stacktrie.go
index b9a360e675..f5d4d7acfa 100644
--- a/trie/stacktrie.go
+++ b/trie/stacktrie.go
@@ -27,7 +27,7 @@ import (
 
 var (
 	stPool = sync.Pool{New: func() any { return new(stNode) }}
-	bPool  = sync.Pool{New: func() any { return make([]byte, 0, 32) }}
+	bPool  = newByteslicepool(32, 100)
 	_      = types.TrieHasher((*StackTrie)(nil))
 )
 
@@ -398,7 +398,7 @@ func (t *StackTrie) hash(st *stNode, path []byte) {
 	// 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 {
-		val := bPool.Get().([]byte)
+		val := bPool.Get()
 		val = val[:len(blob)]
 		copy(val, blob)
 		st.val = val
@@ -406,7 +406,7 @@ func (t *StackTrie) hash(st *stNode, path []byte) {
 	}
 	// Write the hash to the 'val'. We allocate a new val here to not mutate
 	// input values.
-	val := bPool.Get().([]byte)
+	val := bPool.Get()
 	val = val[:32]
 	t.h.hashDataTo(blob, val)
 	st.val = val