execution/commitment: arena-allocate prefixNode extensions

execution/commitment/prefix_trie.go

@@ -23,11 +23,14 @@ import (
 )
 
 const prefixSlabSize = 16384
+const prefixExtChunkSize = 64 * 1024
 
 // prefixNode is a path-compressed prefix-trie node keyed on nibbles (each ext byte is one nibble 0x00..0x0F).
 // children is dense: len == popcount(bitmap). subtreeCount is the number of distinct keys in the
 // subtree; re-inserting an existing key merges its update without bumping it.
 type prefixNode struct {
+	// ext is arena-backed: it stays valid only until the owning trie's Reset, which
+	// recycles the chunk in place. A reader that must outlive the batch copies it.
 	ext          []byte
 	children     []*prefixNode
 	plainKey     []byte  // set only where a key terminates
@@ -41,15 +44,20 @@ type prefixSlab struct {
 	nodes [prefixSlabSize]prefixNode
 }
 
-// prefixArena bump-allocates prefixNodes from a list of slabs.
+// prefixArena bump-allocates prefixNodes from a list of slabs and leaf extensions from a list of byte chunks.
 type prefixArena struct {
-	slabs   []*prefixSlab
-	slabIdx int
-	nextIdx int
+	slabs       []*prefixSlab
+	slabIdx     int
+	nextIdx     int
+	extChunks   [][]byte
+	extChunkIdx int
 }
 
 func newPrefixArena() *prefixArena {
-	return &prefixArena{slabs: []*prefixSlab{new(prefixSlab)}}
+	return &prefixArena{
+		slabs:     []*prefixSlab{new(prefixSlab)},
+		extChunks: [][]byte{make([]byte, 0, prefixExtChunkSize)},
+	}
 }
 
 func (a *prefixArena) allocNode() *prefixNode {
@@ -66,7 +74,34 @@ func (a *prefixArena) allocNode() *prefixNode {
 	return n
 }
 
-// resetArena clears touched nodes for reuse, keeping the first slab and releasing the rest.
+// allocExt copies b into the current chunk, swapping in a fresh chunk instead of growing this one
+// so sub-slices already handed out keep their backing array. An extension larger than a chunk gets
+// its own allocation rather than forcing a chunk to grow under live sub-slices.
+func (a *prefixArena) allocExt(b []byte) []byte {
+	if len(b) == 0 {
+		return nil
+	}
+	if len(b) > prefixExtChunkSize {
+		own := make([]byte, len(b))
+		copy(own, b)
+		return own
+	}
+	chunk := a.extChunks[a.extChunkIdx]
+	if cap(chunk)-len(chunk) < len(b) {
+		a.extChunkIdx++
+		if a.extChunkIdx >= len(a.extChunks) {
+			a.extChunks = append(a.extChunks, make([]byte, 0, prefixExtChunkSize))
+		}
+		chunk = a.extChunks[a.extChunkIdx]
+	}
+	off := len(chunk)
+	chunk = append(chunk, b...)
+	a.extChunks[a.extChunkIdx] = chunk
+	return chunk[off:len(chunk):len(chunk)]
+}
+
+// resetArena clears touched nodes for reuse, keeping the first slab and releasing the rest, and
+// truncates the extension chunks in place so the next batch refills them without reallocating.
 func (a *prefixArena) resetArena() {
 	for i := 0; i <= a.slabIdx && i < len(a.slabs); i++ {
 		limit := prefixSlabSize
@@ -80,6 +115,11 @@ func (a *prefixArena) resetArena() {
 	a.slabs = a.slabs[:1]
 	a.slabIdx = 0
 	a.nextIdx = 0
+
+	for i := range a.extChunks {
+		a.extChunks[i] = a.extChunks[i][:0]
+	}
+	a.extChunkIdx = 0
 }
 
 func (a *prefixArena) nodeCount() int {
@@ -149,7 +189,7 @@ func (t *prefixTrie) Insert(hashedKey, plainKey []byte, update *Update) (isNew b
 			node.plainKey = nil
 			node.update = nil
 
-			node.ext = oldExt[:m]
+			node.ext = oldExt[:m:m]
 
 			if m == len(remain) {
 				// Key ends inside the old extension: one child, no new sibling.
@@ -163,7 +203,7 @@ func (t *prefixTrie) Insert(hashedKey, plainKey []byte, update *Update) (isNew b
 
 			newLeaf := t.arena.allocNode()
 			newNib := remain[m]
-			newLeaf.ext = append([]byte(nil), remain[m+1:]...)
+			newLeaf.ext = t.arena.allocExt(remain[m+1:])
 			newLeaf.subtreeCount = 1
 			newLeaf.plainKey = plainKey
 			newLeaf.update = update
@@ -204,7 +244,7 @@ func (t *prefixTrie) Insert(hashedKey, plainKey []byte, update *Update) (isNew b
 		idx, ok := childIndex(node, nib)
 		if !ok {
 			newLeaf := t.arena.allocNode()
-			newLeaf.ext = append([]byte(nil), hashedKey[keyOffset+1:]...)
+			newLeaf.ext = t.arena.allocExt(hashedKey[keyOffset+1:])
 			newLeaf.subtreeCount = 1
 			newLeaf.plainKey = plainKey
 			newLeaf.update = update

execution/commitment/prefix_trie_bench_test.go

@@ -0,0 +1,53 @@
+// Copyright 2026 The Erigon Authors
+// This file is part of Erigon.
+//
+// Erigon is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// Erigon is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with Erigon. If not, see <http://www.gnu.org/licenses/>.
+
+package commitment
+
+import (
+	"fmt"
+	"testing"
+)
+
+// The Benchmark_Commitment_* family cannot measure the trie build: runParallelBench
+// calls WrapKeyUpdates, which drives every Insert, outside the timed region. This
+// times the build itself across Reset cycles, so the arena's per-batch reuse is
+// what the numbers reflect.
+func Benchmark_PrefixTrieBuildAcrossResets(b *testing.B) {
+	for _, keys := range []int{5_000, 20_000} {
+		b.Run(fmt.Sprintf("%dk-keys", keys/1000), func(b *testing.B) {
+			const keyLen = 64
+			corpus := make([][]byte, keys)
+			for i := range corpus {
+				k := make([]byte, keyLen)
+				v := i
+				for j := range keyLen {
+					k[j] = byte(v % 16)
+					v /= 3
+				}
+				corpus[i] = k
+			}
+			tr := newPrefixTrie()
+			b.ReportAllocs()
+			b.ResetTimer()
+			for range b.N {
+				for _, k := range corpus {
+					tr.Insert(k, k, nil)
+				}
+				tr.Reset()
+			}
+		})
+	}
+}

execution/commitment/prefix_trie_test.go

@@ -225,6 +225,81 @@ func TestPrefixTrieArenaReuse(t *testing.T) {
 	assert.Equal(t, first, tr.arena.nodeCount())
 }
 
+func TestPrefixArenaAllocExt(t *testing.T) {
+	t.Run("capEqualsLen", func(t *testing.T) {
+		a := newPrefixArena()
+		ext := a.allocExt([]byte{1, 2, 3, 4, 5})
+		assert.Equal(t, []byte{1, 2, 3, 4, 5}, ext)
+		assert.Equal(t, len(ext), cap(ext), "allocExt must not hand out spare capacity")
+	})
+
+	t.Run("emptyInputReturnsNil", func(t *testing.T) {
+		a := newPrefixArena()
+		assert.Nil(t, a.allocExt(nil))
+		assert.Nil(t, a.allocExt([]byte{}))
+	})
+
+	t.Run("appendPastCapDoesNotAliasNextExtension", func(t *testing.T) {
+		a := newPrefixArena()
+		first := a.allocExt([]byte{1, 2, 3})
+		second := a.allocExt([]byte{4, 5, 6})
+
+		first = append(first, 0xFF, 0xFF, 0xFF, 0xFF)
+
+		assert.Equal(t, []byte{4, 5, 6}, second, "growing one extension past its cap must not corrupt the next")
+		assert.Equal(t, []byte{1, 2, 3, 0xFF, 0xFF, 0xFF, 0xFF}, first)
+	})
+}
+
+func TestPrefixTrieExtSurvivesChunkBoundary(t *testing.T) {
+	tr := newPrefixTrie()
+
+	// keys share nibbles [4:keyLen) so leaf extensions stay long; nibbles [0:4) alone already
+	// make every key distinct, which is what drives enough allocExt traffic to cross a chunk.
+	const keyLen = 32
+	const total = 2 * prefixExtChunkSize / keyLen
+	want := make(map[string]bool, total)
+	for i := range total {
+		k := make([]byte, keyLen)
+		v := i
+		for j := range 4 {
+			k[j] = byte(v % 16)
+			v /= 16
+		}
+		want[string(k)] = true
+		tr.Insert(k, nil, nil)
+	}
+	require.Greater(t, len(tr.arena.extChunks), 1, "test must actually cross a chunk boundary")
+
+	got := make(map[string]bool, total)
+	for _, e := range collectWalk(tr) {
+		if len(e.prefix) == keyLen {
+			got[string(e.prefix)] = true
+		}
+	}
+	assert.Equal(t, want, got, "leaf extensions must reproduce their original key bytes after crossing a chunk boundary")
+}
+
+func TestPrefixTrieArenaReusesExtChunkBacking(t *testing.T) {
+	tr := newPrefixTrie()
+	tr.Insert(nibs(0x01, 0x02, 0x03, 0x04), nil, nil)
+
+	chunk := tr.arena.extChunks[0]
+	full := chunk[:cap(chunk)]
+
+	tr.Reset()
+
+	require.Len(t, tr.arena.extChunks, 1, "Reset must trim trailing chunks")
+	assert.Empty(t, tr.arena.extChunks[0], "Reset must truncate the reused chunk's length")
+	assert.Equal(t, cap(chunk), cap(tr.arena.extChunks[0]), "Reset must keep the chunk's capacity")
+
+	tr.Insert(nibs(0x05, 0x06, 0x07, 0x08), nil, nil)
+	require.Len(t, tr.root.children, 1)
+	newExt := tr.root.children[0].ext
+	assert.Equal(t, nibs(0x06, 0x07, 0x08), newExt)
+	assert.Equal(t, newExt, full[:len(newExt)], "post-reset extension must land in the same backing array as before Reset")
+}
+
 func TestPrefixTrieArenaSpansMultipleSlabs(t *testing.T) {
 	tr := newPrefixTrie()
 	// Allocate directly to cross the slab boundary; reaching it via inserts needs >prefixSlabSize keys.
@@ -431,3 +506,42 @@ func TestPrefixTrieInsertDuplicateMerges(t *testing.T) {
 	// Merge is copy-on-write: a concurrent fold snapshot may still hold the prior update pointer.
 	assert.Equal(t, BalanceUpdate, first.Flags, "merge must not mutate the previously stored update")
 }
+
+func TestPrefixArenaAllocExt_OversizeAndReuse(t *testing.T) {
+	t.Run("extension larger than a chunk gets its own backing", func(t *testing.T) {
+		a := newPrefixArena()
+		big := bytes.Repeat([]byte{0x7}, prefixExtChunkSize+1)
+		got := a.allocExt(big)
+		require.Equal(t, big, got)
+		require.Equal(t, len(got), cap(got))
+		require.Len(t, a.extChunks, 1, "an oversize extension must not consume a chunk")
+		require.Empty(t, a.extChunks[0], "the current chunk must be untouched")
+	})
+
+	t.Run("reset refills existing chunks instead of reallocating", func(t *testing.T) {
+		a := newPrefixArena()
+		block := make([]byte, prefixExtChunkSize/4)
+		for range 12 {
+			a.allocExt(block)
+		}
+		require.Greater(t, len(a.extChunks), 1, "test must cross a chunk boundary")
+		grown := len(a.extChunks)
+		backing := make([]*byte, grown)
+		for i := range a.extChunks {
+			backing[i] = &a.extChunks[i][:1][0]
+		}
+
+		a.resetArena()
+		require.Len(t, a.extChunks, grown, "reset must keep the grown chunks")
+		for i := range a.extChunks {
+			require.Empty(t, a.extChunks[i])
+		}
+		for range 12 {
+			a.allocExt(block)
+		}
+		require.Len(t, a.extChunks, grown, "refill must reuse chunks, not allocate new ones")
+		for i := range a.extChunks {
+			require.Same(t, backing[i], &a.extChunks[i][:1][0], "chunk backing array must be reused")
+		}
+	})
+}