execution/commitment: arena-allocate prefixNode extensions

[awskii]

prefixTrie.Insert allocated a fresh []byte for every new leaf's extension. prefixNodes themselves already come from a slab arena that survives across batches; their extension bytes did not, so a block's trie build paid one heap allocation per new key.

Changes

Add a byte-chunk arena alongside the node slabs — 64KB chunks, a new chunk swapped in on overflow rather than grown, so outstanding extensions keep their backing and every returned slice has cap == len. An extension larger than a chunk takes its own exact-capacity allocation instead of forcing a chunk to grow under live sub-slices. resetArena truncates every chunk in place, matching Updates.Reset's arena-ring policy, so a steady-state batch refills them without allocating.

The split path now caps node.ext = oldExt[:m:m]. With per-extension allocations the spare capacity was harmless; once extensions share arena memory, an append there would reach into the bytes oldChild.ext owns.

Benchmark_PrefixTrieBuildAcrossResets (20k keys, insert + Reset per iteration), Apple M5 Max, -benchtime=200x -count=8:

	before	after
sec/op	1.677m	1.448m	−13.69%
B/op	4.365Mi	3.149Mi	−27.85%
allocs/op	39.99k	19.99k	−50.01%

All p=0.000, n=8.

Note on the benchmark choice: Benchmark_Commitment_DirectVsParallel cannot see this change at all — runParallelBench calls WrapKeyUpdates, which drives every Insert, before b.StartTimer(), so the whole trie build sits outside the timer. An earlier revision of this PR quoted figures from that build phase measured through a temporary benchmark; those predated the resetArena fix and are superseded by the table above, which exercises the reset cycle directly.

[Copilot]

Pull request overview

This PR reduces per-key heap allocations during commitment prefix-trie construction by moving leaf-extension storage into an arena, matching the existing slab allocator used for prefixNode objects. This targets the hot path in prefixTrie.Insert where each new leaf previously allocated a fresh []byte for its extension.

Changes:

• Added an extension-byte chunk arena (prefixExtChunkSize chunks) to prefixArena, with allocExt returning slices with cap == len to prevent accidental aliasing via append.
• Updated prefixTrie.Insert to allocate leaf extensions via the new arena instead of append([]byte(nil), ...).
• Added targeted tests to validate extension allocation semantics, chunk-boundary correctness, and reset reuse behavior.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
execution/commitment/prefix_trie.go	Introduces extension-byte chunk arena in prefixArena and routes leaf extension allocation through it to reduce heap allocs.
execution/commitment/prefix_trie_test.go	Adds tests validating allocExt behavior (capacity, nil-on-empty, non-aliasing), chunk-boundary correctness, and reset reuse of chunk backing.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[AskAlexSharov]

Review of the extension arena. No correctness bug that fires today: allocExt copies, the three-index return keeps handed-out slices append-safe, every in-package reader of node.ext copies before the arena is reset, and the background folds read the trie under trieMu with Stop() draining them before Reset(). The notes below are latent hazards, reuse, and one measurement gap.

One item that cannot be anchored to a line in this diff:

The optimized path is not covered by any committed benchmark. runParallelBench (execution/commitment/parallel_streaming_bench_test.go:56) calls WrapKeyUpdates between b.StopTimer() and b.StartTimer(), so the whole prefix-trie build sits outside the timer in every Benchmark_Commitment_*. The PR description names this, but the consequence is that the 27% allocation win is defended by nothing in CI's bench / benchmarks (parallel) job -- a later regression here would be invisible. A build-phase benchmark over the same corpus (or a dedicated variant that keeps WrapKeyUpdates inside the timed region) would close it.

[AskAlexSharov]

extChunkIdx is derivable state: it is always len(a.extChunks)-1. The index only ever advances together with an append to extChunks, and resetArena puts both back to 0 / len 1. So the if a.extChunkIdx >= len(a.extChunks) branch in allocExt is always taken, and the re-fetch after it always yields the freshly appended chunk.

Dropping the field and reading a.extChunks[len(a.extChunks)-1] removes one invariant a reader has to verify by hand. It also matters for the oversize case: the "growing in place is harmless" reasoning holds only while the newly selected chunk is guaranteed empty, which is precisely what this lockstep gives you.

[awskii]

Not applied, and I think this one stops holding once the chunk-retention comment is addressed — the two interact.

extChunkIdx == len(extChunks)-1 was true only because resetArena dropped every chunk but the first, so the arena could never be sitting on a non-last chunk. Now that reset retains the grown chunks truncated, the next batch starts at chunk 0 with several empty chunks above it, and the index is precisely what lets it refill them in order instead of appending past them. Deriving it would make every reset re-grow the tail.

Your underlying point still lands though: the "growing in place is harmless" reasoning did depend on the lockstep. That is no longer what guarantees it — the explicit oversize guard is, so no chunk is ever grown under live sub-slices regardless of which one is current.

[AskAlexSharov]

This is the fourth chunked byte arena in this package, all with the same shape -- 64KB chunks, "the returned sub-slice must stay valid until reset":

• plainKeyArena.intern -- parallel_update.go
• keyArena.copy -- streaming_commitment.go
• Updates.arenaAlloc / arenaEnsureCap -- commitment.go
• prefixArena.allocExt -- here

They differ only in the overflow policy and in whether the returned slice is capped (arenaAlloc returns arena[off:needed], without the three-index form, so its callers do not get the append-safety this PR adds here). One shared byteArena type with a single overflow policy would remove three near-copies and give every caller the capped return. The new constant is also spelled 65536 while the other three use 64 * 1024.

[awskii]

Agreed on all four counts, and the constant is now 64 * 1024 to match its siblings.

Not unifying them in this PR: keyArena and plainKeyArena are both moving under the streaming-commitment removal (#23191), which relocates keyArena into streaming_deep_fold.go, and Updates.arenaAlloc has the different contract you point out — uncapped return, ring of arenas indexed by generation. Folding all four together across those in-flight moves would make both diffs hard to review. Worth a follow-up once #23191 and #23183 land, and the capped return is the part actually worth propagating — arenaAlloc's callers not getting append-safety is a live sharp edge, not just duplication.

[AskAlexSharov]

Worth noting for the follow-up: the per-leaf extension was one of two heap allocations per new node. The other is the children slice -- []*prefixNode{oldChild} and []*prefixNode{oldChild, newLeaf} in the split branch, plus the append(node.children, nil) growth just below -- and merged := &Update{} on the duplicate-key path. That is consistent with the residual 1.326M allocs/op after this change. A children region in the same arena (or a fixed [16]*prefixNode block per node) would compound with this PR rather than overlap it.

[awskii]

Noted, and it matches what the profile shows — the residual is dominated by the children slices, with merged := &Update{} behind them.

Deliberately out of scope here: the children region changes node layout rather than just where bytes come from, so it wants its own before/after. A fixed [16]*prefixNode per node trades ~128 bytes per node against the current dense slice, which is the opposite trade for a trie whose nodes are mostly sparse — a children arena with the dense layout kept is the version I would measure first.

[awskii]

On the unanchored item — the measurement gap is real and now closed.

Benchmark_PrefixTrieBuildAcrossResets is committed (prefix_trie_bench_test.go). You are right that the win was defended by nothing: I had measured through a throwaway benchmark and deleted it, so CI's bench / benchmarks (parallel) job would not have caught a regression here. It times the build across Reset cycles rather than a single build, which is also what exercises the chunk reuse — and it is the benchmark the numbers in the description now come from, so they are reproducible.

That choice turned out to matter beyond coverage: your resetArena note was right that "a steady-state batch allocates no extension bytes" only held below one chunk, and it was the reset-cycle benchmark that showed it. With all chunks truncated in place instead of dropped:

sec/op     1.677m -> 1.448m  -13.69%
B/op      4.365Mi -> 3.149Mi  -27.85%
allocs/op  39.99k ->  19.99k  -50.01%

p=0.000, n=8. The earlier 27%/4.7% figures were from the pre-fix policy measured through the throwaway benchmark; they are gone from the description.

Everything else in the review is applied except the extChunkIdx removal, which stops holding once chunks are retained — replied inline. The four-arena unification I would rather do after #23191 lands, since it moves keyArena out of streaming_commitment.go entirely.