execution/commitment: let trie warmup borrow branch bytes instead of copying them

execution/commitment/commitment.go

@@ -121,6 +121,14 @@ type PatriciaContext interface {
 	Storage(plainKey []byte) (*Update, error)
 }
 
+// BranchNoCopyReader reads a branch without copying it. The result aliases the
+// reader's memory: it is valid until the next call on the same context, and must
+// not be mutated. Branch copies because the trie's consumers keep branch data
+// past that.
+type BranchNoCopyReader interface {
+	BranchNoCopy(prefix []byte) ([]byte, kv.Step, error)
+}
+
 type TrieVariant string
 
 const (

execution/commitment/commitmentdb/commitment_context.go

@@ -954,7 +954,7 @@ func NewTrieContextRo(reader StateReader, stepSize uint64) *TrieContext {
 }
 
 func (sdc *TrieContext) Branch(pref []byte) ([]byte, kv.Step, error) {
-	enc, step, err := sdc.readDomain(kv.CommitmentDomain, pref)
+	enc, step, err := sdc.branch(pref)
 	if err != nil {
 		return nil, 0, err
 	}
@@ -963,10 +963,26 @@ func (sdc *TrieContext) Branch(pref []byte) ([]byte, kv.Step, error) {
 	// underlying state cache / getter aliases shared storage that another goroutine
 	// (concurrent commitment workers) can recycle. Own the bytes at the trie-context
 	// boundary so all downstream consumers are safe.
+	return bytes.Clone(enc), step, nil
+}
+
+// BranchNoCopy returns the branch bytes uncopied, for a caller that consumes them
+// before its next read -- the trie warmup. The result is valid only until the
+// next call on this context and must not be mutated; anything that keeps branch
+// data uses Branch.
+func (sdc *TrieContext) BranchNoCopy(pref []byte) ([]byte, kv.Step, error) {
+	return sdc.branch(pref)
+}
+
+func (sdc *TrieContext) branch(pref []byte) ([]byte, kv.Step, error) {
+	enc, step, err := sdc.readDomain(kv.CommitmentDomain, pref)
+	if err != nil {
+		return nil, 0, err
+	}
 	if sdc.traceW != nil {
 		fmt.Fprintf(sdc.traceW, "[SDC] Branch read %x => %x\n", pref, enc)
 	}
-	return bytes.Clone(enc), step, nil
+	return enc, step, nil
 }
 
 func (sdc *TrieContext) PutBranch(prefix []byte, data []byte, prevData []byte) error {

execution/commitment/warmuper.go

@@ -29,6 +29,7 @@ import (
 
 	"github.com/erigontech/erigon/common"
 	"github.com/erigontech/erigon/common/log/v3"
+	"github.com/erigontech/erigon/db/kv"
 	"github.com/erigontech/erigon/execution/commitment/nibbles"
 )
 
@@ -115,6 +116,11 @@ func (w *Warmuper) Start() {
 				return errors.New("warmup trie context factory returned nil PatriciaContext")
 			}
 
+			readBranch := trieCtx.Branch
+			if w, ok := trieCtx.(BranchNoCopyReader); ok {
+				readBranch = w.BranchNoCopy
+			}
+
 			for {
 				select {
 				case <-w.ctx.Done():
@@ -123,7 +129,7 @@ func (w *Warmuper) Start() {
 					if !ok {
 						return nil
 					}
-					w.warmupKey(trieCtx, item.hashedKey, item.startDepth)
+					w.warmupKey(readBranch, item.hashedKey, item.startDepth)
 					w.keysProcessed.Add(1)
 					w.releaseGen(item.gen)
 				}
@@ -141,13 +147,15 @@ func (w *Warmuper) Start() {
 	})
 }
 
-func (w *Warmuper) warmupKey(trieCtx PatriciaContext, hashedKey []byte, startDepth int) {
+// warmupKey descends one key, parsing each branch only to pick the next nibble.
+// It keeps nothing, so the caller can pass the non-copying read.
+func (w *Warmuper) warmupKey(readBranch func([]byte) ([]byte, kv.Step, error), hashedKey []byte, startDepth int) {
 	depth := startDepth
 	var compactBuf [maxCompactKeyLen]byte
 	for depth <= len(hashedKey) && depth <= w.maxDepth {
 		prefix := nibbles.HexToCompactInto(compactBuf[:], hashedKey[:depth])
 
-		branchData, _, err := trieCtx.Branch(prefix)
+		branchData, _, err := readBranch(prefix)
 		if err != nil {
 			log.Debug(fmt.Sprintf("[%s][warmup] failed to get branch", w.logPrefix),
 				"prefix", common.Bytes2Hex(prefix), "error", err)

execution/commitment/warmuper_test.go

@@ -17,10 +17,16 @@
 package commitment
 
 import (
+	"bytes"
 	"context"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/erigontech/erigon/db/kv"
 )
 
 func TestWarmuperFactoryMustNotOutliveCloseAndWait(t *testing.T) {
@@ -255,3 +261,55 @@ func TestCloseLeavesWorkChannelOpen(t *testing.T) {
 	default:
 	}
 }
+
+// countingBranchCtx records which read path the warmuper took, and holds
+// BranchNoCopy to its contract: each call poisons what the previous one returned,
+// so a caller that kept the bytes reads garbage and descends differently.
+type countingBranchCtx struct {
+	src      []byte
+	last     []byte
+	owned    atomic.Int64
+	borrowed atomic.Int64
+}
+
+func (c *countingBranchCtx) Branch(prefix []byte) ([]byte, kv.Step, error) {
+	c.owned.Add(1)
+	return bytes.Clone(c.src), 0, nil
+}
+
+func (c *countingBranchCtx) BranchNoCopy(prefix []byte) ([]byte, kv.Step, error) {
+	c.borrowed.Add(1)
+	for i := range c.last {
+		c.last[i] = 0xff
+	}
+	c.last = bytes.Clone(c.src)
+	return c.last, 0, nil
+}
+
+func (c *countingBranchCtx) PutBranch(prefix, data, prevData []byte) error { return nil }
+func (c *countingBranchCtx) Account(plainKey []byte) (*Update, error)      { return nil, nil }
+func (c *countingBranchCtx) Storage(plainKey []byte) (*Update, error)      { return nil, nil }
+
+// TestWarmuperUsesBranchNoCopy pins the warmuper on the non-copying read. It keeps
+// nothing it reads, and it issues several times more branch reads than the fold
+// does, so an owning read there is the bulk of the copying.
+func TestWarmuperUsesBranchNoCopy(t *testing.T) {
+	t.Parallel()
+	// touchMap, afterMap with nibble 0 set, then one cell with no fields.
+	ctx := &countingBranchCtx{src: []byte{0x00, 0x01, 0x00, 0x01, 0x00}}
+	w := NewWarmuper(context.Background(), WarmupConfig{
+		Enabled:    true,
+		CtxFactory: func(context.Context) (PatriciaContext, func()) { return ctx, nil },
+		NumWorkers: 1,
+		MaxDepth:   WarmupMaxDepth,
+	})
+	w.Start()
+	w.WarmKey([]byte{0, 0, 0, 0}, 0, 0)
+	require.NoError(t, w.WaitBufferFree(0))
+	w.CloseAndWait()
+
+	require.Zero(t, ctx.owned.Load(), "warmup still copies branch bytes it drops immediately")
+	// One read per nibble of the key, then one more that stops at its end. A
+	// caller holding a poisoned branch would read 0xff and stop short.
+	require.Equal(t, int64(5), ctx.borrowed.Load(), "warmup descent read the wrong number of branches")
+}