From 4eb7032c97358b4a3073e74a27f6b07128e87896 Mon Sep 17 00:00:00 2001
From: Martin Holst Swende <martin@swende.se>
Date: Mon, 11 Nov 2024 14:20:22 +0100
Subject: [PATCH] trie: stacktrie allocation reduction via key scratchspace

---
 trie/encoding.go  | 11 +++++++++++
 trie/stacktrie.go | 18 ++++++++++++++++--
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/trie/encoding.go b/trie/encoding.go
index 3284d3f8f0..d795ba1895 100644
--- a/trie/encoding.go
+++ b/trie/encoding.go
@@ -104,6 +104,17 @@ func keybytesToHex(str []byte) []byte {
 	return nibbles
 }
 
+// writeHexKey writes the hexkey into the given slice.
+// OBS! This method omits the termination flag.
+// OBS! The dst slice must be at least 2x as large as the key
+func writeHexKey(dst []byte, key []byte) {
+	_ = dst[2*len(key)-1]
+	for i, b := range key {
+		dst[i*2] = b / 16
+		dst[i*2+1] = b % 16
+	}
+}
+
 // hexToKeybytes turns hex nibbles into key bytes.
 // This can only be used for keys of even length.
 func hexToKeybytes(hex []byte) []byte {
diff --git a/trie/stacktrie.go b/trie/stacktrie.go
index d66bba3f11..409e8a1136 100644
--- a/trie/stacktrie.go
+++ b/trie/stacktrie.go
@@ -47,6 +47,8 @@ type StackTrie struct {
 	h          *hasher
 	last       []byte
 	onTrieNode OnTrieNode
+
+	keyScratch []byte
 }
 
 // NewStackTrie allocates and initializes an empty trie. The committed nodes
@@ -64,7 +66,18 @@ func (t *StackTrie) Update(key, value []byte) error {
 	if len(value) == 0 {
 		return errors.New("trying to insert empty (deletion)")
 	}
-	k := t.TrieKey(key)
+	var k []byte
+	{
+		// We can reuse the key scratch area, but only if the insert-method
+		// never holds on to it.
+		if cap(t.keyScratch) < 2*len(key) { // realloc to ensure sufficient cap
+			t.keyScratch = make([]byte, 2*len(key), 2*len(key))
+		}
+		// resize to ensure correct size
+		t.keyScratch = t.keyScratch[:2*len(key)]
+		writeHexKey(t.keyScratch, key)
+		k = t.keyScratch
+	}
 	if bytes.Compare(t.last, k) >= 0 {
 		return errors.New("non-ascending key order")
 	}
@@ -152,6 +165,7 @@ func (n *stNode) getDiffIndex(key []byte) int {
 
 // Helper function to that inserts a (key, value) pair into
 // the trie.
+// The key is not retained by this method, but always copied if needed.
 func (t *StackTrie) insert(st *stNode, key, value []byte, path []byte) {
 	switch st.typ {
 	case branchNode: /* Branch */
@@ -283,7 +297,7 @@ func (t *StackTrie) insert(st *stNode, key, value []byte, path []byte) {
 
 	case emptyNode: /* Empty */
 		st.typ = leafNode
-		st.key = key
+		st.key = append(st.key, key...)
 		st.val = value
 
 	case hashedNode: