Skip to content

replaceShortenedKeysInBranch: less re-allocs #19584

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
AskAlexSharov merged 7 commits into from
Mar 4, 2026
Merged

replaceShortenedKeysInBranch: less re-allocs #19584

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
642b78a
save
AskAlexSharov Mar 3, 2026
ca68093
save
AskAlexSharov Mar 3, 2026
92fefb5
save
AskAlexSharov Mar 4, 2026
3a909c7
Merge remote-tracking branch 'origin/main' into alex/shortened_key_bu…
AskAlexSharov Mar 4, 2026
7528af1
save
AskAlexSharov Mar 4, 2026
471e1da
save
AskAlexSharov Mar 4, 2026
5d12fd2
save
AskAlexSharov Mar 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 16 additions & 6 deletions db/integrity/commitment_integirty.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -351,7 +351,7 @@ func checkDerefBranch(
branchData := commitment.BranchData(branchValue)
var integrityErr error

newBranchData, _ = branchData.ReplacePlainKeys(newBranchValueBuf[:0], func(key []byte, isStorage bool) ([]byte, error) {
newBranchData, newBranchValueBuf, _ = branchData.ReplacePlainKeys(newBranchValueBuf[:0], func(key []byte, isStorage bool) ([]byte, error) {
if trace {
logger.Trace(
"checking commitment deref for branch",
Expand Down Expand Up @@ -967,6 +967,7 @@ func checkStateCorrespondenceBase(ctx context.Context, file state.VisibleFile, s
branchValueBuf := make([]byte, 0, datasize.MB.Bytes())
var branchKeys uint64
var integrityErr error
var branchDerefBuf []byte

for i := 0; commReader.HasNext(); i++ {
if i%1024 == 0 {
Expand Down Expand Up @@ -1048,7 +1049,8 @@ func checkStateCorrespondenceBase(ctx context.Context, file state.VisibleFile, s
}
// The callback returns nil (keep original key in output) because the result of ReplacePlainKeys
// is discarded (_): all side-effects happen before the return.
_, err := branchData.ReplacePlainKeys(nil, func(key []byte, isStorage bool) ([]byte, error) {
var err error
_, branchDerefBuf, err = branchData.ReplacePlainKeys(branchDerefBuf[:0], func(key []byte, isStorage bool) ([]byte, error) {
var checkErr error
if isStorage {
checkErr = checkKey(key, stoCollector, storageReader, length.Addr+length.Hash, "storage")
Expand Down Expand Up @@ -1183,6 +1185,7 @@ func checkStateCorrespondenceReverse(ctx context.Context, file state.VisibleFile
var branchKeys uint64
var integrityErr error
var extractedAccKeys, extractedStoKeys, skippedAccKeys, skippedStoKeys uint64
var branchDerefBuf []byte

// Phase 1: Walk commitment branches, collect extracted plain keys into ETL collectors.
for i := 0; commReader.HasNext(); i++ {
Expand Down Expand Up @@ -1226,7 +1229,8 @@ func checkStateCorrespondenceReverse(ctx context.Context, file state.VisibleFile
continue
}

_, err := branchData.ReplacePlainKeys(nil, func(key []byte, isStorage bool) ([]byte, error) {
var err error
_, branchDerefBuf, err = branchData.ReplacePlainKeys(branchDerefBuf[:0], func(key []byte, isStorage bool) ([]byte, error) {
if isStorage {
plainKey := key
if len(key) != length.Addr+length.Hash {
Expand Down Expand Up @@ -1565,6 +1569,7 @@ func extractCommitmentRefsToCollectors(ctx context.Context, file state.VisibleFi
branchKeyBuf := make([]byte, 0, 128)
branchValueBuf := make([]byte, 0, datasize.MB.Bytes())
plainKeyBuf := make([]byte, 0, length.Addr+length.Hash)
var branchDerefBuf []byte

for commReader.HasNext() {
select {
Expand All @@ -1588,7 +1593,8 @@ func extractCommitmentRefsToCollectors(ctx context.Context, file state.VisibleFi
continue
}

if _, err := branchData.ReplacePlainKeys(nil, func(key []byte, isStorage bool) ([]byte, error) { //nolint:gocritic
var err error
_, branchDerefBuf, err = branchData.ReplacePlainKeys(branchDerefBuf[:0], func(key []byte, isStorage bool) ([]byte, error) { //nolint:gocritic
if isStorage {
plainKey := key
if len(key) != length.Addr+length.Hash {
Expand Down Expand Up @@ -1621,7 +1627,8 @@ func extractCommitmentRefsToCollectors(ctx context.Context, file state.VisibleFi
}
}
return plainKey, nil
}); err != nil {
})
if err != nil {
return err
}
}
Expand Down Expand Up @@ -1709,6 +1716,7 @@ func checkHashVerification(ctx context.Context, file state.VisibleFile, stepSize

plainKeyBuf := make([]byte, 0, length.Addr+length.Hash)
valBuf := make([]byte, 0, 128)
var branchDerefBuf []byte

for item := range workCh {
select {
Expand All @@ -1726,7 +1734,9 @@ func checkHashVerification(ctx context.Context, file state.VisibleFile, stepSize

// We need branch data with plain keys for VerifyBranchHashes.
// Walk the branch to extract + resolve all keys and read values.
resolvedBranchData, err := branchData.ReplacePlainKeys(nil, func(key []byte, isStorage bool) ([]byte, error) {
var resolvedBranchData commitment.BranchData
var err error
resolvedBranchData, branchDerefBuf, err = branchData.ReplacePlainKeys(branchDerefBuf[:0], func(key []byte, isStorage bool) ([]byte, error) {
if isStorage {
plainKey := key
isRef := len(key) != length.Addr+length.Hash
Expand Down
2 changes: 2 additions & 0 deletions db/state/aggregator.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1754,6 +1754,8 @@ type AggregatorRoTx struct {
d [kv.DomainLen]*DomainRoTx
iis []*InvertedIndexRoTx

branchDerefBuf []byte // scratch buffer reused across replaceShortenedKeysInBranch calls

_leakID uint64 // set only if TRACE_AGG=true
}

Expand Down
14 changes: 9 additions & 5 deletions db/state/domain_committed.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,8 +85,8 @@ func (at *AggregatorRoTx) replaceShortenedKeysInBranch(prefix []byte, branch com
}
}

aux := make([]byte, 0, 256)
return branch.ReplacePlainKeys(aux, func(key []byte, isStorage bool) ([]byte, error) {
var result commitment.BranchData
result, at.branchDerefBuf, err = branch.ReplacePlainKeys(at.branchDerefBuf[:0], func(key []byte, isStorage bool) ([]byte, error) {
if isStorage {
if len(key) == length.Addr+length.Hash {
return nil, nil // save storage key as is
Expand Down Expand Up @@ -121,6 +121,10 @@ func (at *AggregatorRoTx) replaceShortenedKeysInBranch(prefix []byte, branch com
}
return apkBuf, nil
})
if err != nil {
return nil, err
}
return result, nil
}

func DecodeReferenceKey(from []byte) uint64 {
Expand Down Expand Up @@ -427,12 +431,12 @@ func (dt *DomainRoTx) commitmentValTransformDomain(rng MergeRange, accounts, sto
return shortened, nil
}

temp, err := commitment.BranchData(valBuf).ReplacePlainKeys(dt.comBuf[:0], replacer)
var branchData commitment.BranchData
branchData, dt.comBuf, err = commitment.BranchData(valBuf).ReplacePlainKeys(dt.comBuf[:0], replacer)
if err != nil {
return nil, err
}
dt.comBuf = append(dt.comBuf[:0], temp...) // cover branch data case
return dt.comBuf, nil
return branchData, nil
}

return vt, nil
Expand Down
79 changes: 41 additions & 38 deletions execution/commitment/commitment.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -871,137 +871,140 @@ func (branchData BranchData) String() string {
return sb.String()
}

// if fn returns nil, the original key will be copied from branchData
func (branchData BranchData) ReplacePlainKeys(newData []byte, fn func(key []byte, isStorage bool) (newKey []byte, err error)) (BranchData, error) {
// if fn returns nil, the original key will be copied from branchData.
// Returns (result, buf, err): result may point into mmap when returned unchanged;
// buf is always the heap scratch buffer (grown from the input buf), safe to retain
// across calls. Callers that need a heap-owned copy of result should bytes.Clone(result).
func (branchData BranchData) ReplacePlainKeys(buf []byte, fn func(key []byte, isStorage bool) (newKey []byte, err error)) (BranchData, []byte, error) {
if len(branchData) < 4 {
return branchData, nil
return branchData, buf, nil
}

var numBuf [binary.MaxVarintLen64]byte
touchMap := binary.BigEndian.Uint16(branchData[0:])
afterMap := binary.BigEndian.Uint16(branchData[2:])
if touchMap&afterMap == 0 {
return branchData, nil
return branchData, buf, nil
}
pos := 4
newData = append(newData[:0], branchData[:4]...)
buf = append(buf[:0], branchData[:4]...)
for bitset, j := touchMap&afterMap, 0; bitset != 0; j++ {
bit := bitset & -bitset
fields := cellFields(branchData[pos])
newData = append(newData, byte(fields))
buf = append(buf, byte(fields))
pos++
if fields&fieldExtension != 0 {
l, n := binary.Uvarint(branchData[pos:])
if n == 0 {
return nil, errors.New("replacePlainKeys buffer too small for hashedKey len")
return nil, buf, errors.New("replacePlainKeys buffer too small for hashedKey len")
} else if n < 0 {
return nil, errors.New("replacePlainKeys value overflow for hashedKey len")
return nil, buf, errors.New("replacePlainKeys value overflow for hashedKey len")
}
newData = append(newData, branchData[pos:pos+n]...)
buf = append(buf, branchData[pos:pos+n]...)
pos += n
if len(branchData) < pos+int(l) {
return nil, fmt.Errorf("replacePlainKeys buffer too small for hashedKey: expected %d got %d", pos+int(l), len(branchData))
return nil, buf, fmt.Errorf("replacePlainKeys buffer too small for hashedKey: expected %d got %d", pos+int(l), len(branchData))
}
if l > 0 {
newData = append(newData, branchData[pos:pos+int(l)]...)
buf = append(buf, branchData[pos:pos+int(l)]...)
pos += int(l)
}
}
if fields&fieldAccountAddr != 0 {
l, n := binary.Uvarint(branchData[pos:])
if n == 0 {
return nil, errors.New("replacePlainKeys buffer too small for accountAddr len")
return nil, buf, errors.New("replacePlainKeys buffer too small for accountAddr len")
} else if n < 0 {
return nil, errors.New("replacePlainKeys value overflow for accountAddr len")
return nil, buf, errors.New("replacePlainKeys value overflow for accountAddr len")
}
pos += n
if len(branchData) < pos+int(l) {
return nil, fmt.Errorf("replacePlainKeys buffer too small for accountAddr: expected %d got %d", pos+int(l), len(branchData))
return nil, buf, fmt.Errorf("replacePlainKeys buffer too small for accountAddr: expected %d got %d", pos+int(l), len(branchData))
}
if l > 0 {
pos += int(l)
}
newKey, err := fn(branchData[pos-int(l):pos], false)
if err != nil {
return nil, err
return nil, buf, err
}
if newKey == nil {
// invariant: fn returns nil (keep original) only for plain addr keys (length.Addr bytes)
newData = append(newData, branchData[pos-int(l)-n:pos]...)
buf = append(buf, branchData[pos-int(l)-n:pos]...)
} else {
// invariant: newKey is a short reference (≤8 bytes) or full plain addr (length.Addr bytes)
n = binary.PutUvarint(numBuf[:], uint64(len(newKey)))
newData = append(newData, numBuf[:n]...)
newData = append(newData, newKey...)
buf = append(buf, numBuf[:n]...)
buf = append(buf, newKey...)
}
}
if fields&fieldStorageAddr != 0 {
l, n := binary.Uvarint(branchData[pos:])
if n == 0 {
return nil, errors.New("replacePlainKeys buffer too small for storageAddr len")
return nil, buf, errors.New("replacePlainKeys buffer too small for storageAddr len")
} else if n < 0 {
return nil, errors.New("replacePlainKeys value overflow for storageAddr len")
return nil, buf, errors.New("replacePlainKeys value overflow for storageAddr len")
}
pos += n
if len(branchData) < pos+int(l) {
return nil, fmt.Errorf("replacePlainKeys buffer too small for storageAddr: expected %d got %d", pos+int(l), len(branchData))
return nil, buf, fmt.Errorf("replacePlainKeys buffer too small for storageAddr: expected %d got %d", pos+int(l), len(branchData))
}
if l > 0 {
pos += int(l)
}
newKey, err := fn(branchData[pos-int(l):pos], true)
if err != nil {
return nil, err
return nil, buf, err
}
if newKey == nil {
// invariant: fn returns nil (keep original) only for plain storage keys (length.Addr+length.Hash bytes)
newData = append(newData, branchData[pos-int(l)-n:pos]...) // -n to include length
buf = append(buf, branchData[pos-int(l)-n:pos]...) // -n to include length
} else {
// invariant: newKey is a short reference (≤8 bytes) or full plain storage key (length.Addr+length.Hash bytes)
n = binary.PutUvarint(numBuf[:], uint64(len(newKey)))
newData = append(newData, numBuf[:n]...)
newData = append(newData, newKey...)
buf = append(buf, numBuf[:n]...)
buf = append(buf, newKey...)
}
}
if fields&fieldHash != 0 {
l, n := binary.Uvarint(branchData[pos:])
if n == 0 {
return nil, errors.New("replacePlainKeys buffer too small for hash len")
return nil, buf, errors.New("replacePlainKeys buffer too small for hash len")
} else if n < 0 {
return nil, errors.New("replacePlainKeys value overflow for hash len")
return nil, buf, errors.New("replacePlainKeys value overflow for hash len")
}
newData = append(newData, branchData[pos:pos+n]...)
buf = append(buf, branchData[pos:pos+n]...)
pos += n
if len(branchData) < pos+int(l) {
return nil, fmt.Errorf("replacePlainKeys buffer too small for hash: expected %d got %d", pos+int(l), len(branchData))
return nil, buf, fmt.Errorf("replacePlainKeys buffer too small for hash: expected %d got %d", pos+int(l), len(branchData))
}
if l > 0 {
newData = append(newData, branchData[pos:pos+int(l)]...)
buf = append(buf, branchData[pos:pos+int(l)]...)
pos += int(l)
}
}
if fields&fieldStateHash != 0 {
l, n := binary.Uvarint(branchData[pos:])
if n == 0 {
return nil, errors.New("replacePlainKeys buffer too small for acLeaf hash len")
return nil, buf, errors.New("replacePlainKeys buffer too small for acLeaf hash len")
} else if n < 0 {
return nil, errors.New("replacePlainKeys value overflow for acLeafhash len")
return nil, buf, errors.New("replacePlainKeys value overflow for acLeafhash len")
}
newData = append(newData, branchData[pos:pos+n]...)
buf = append(buf, branchData[pos:pos+n]...)
pos += n
if len(branchData) < pos+int(l) {
return nil, fmt.Errorf("replacePlainKeys buffer too small for LeafHash: expected %d got %d", pos+int(l), len(branchData))
return nil, buf, fmt.Errorf("replacePlainKeys buffer too small for LeafHash: expected %d got %d", pos+int(l), len(branchData))
}
if l > 0 {
newData = append(newData, branchData[pos:pos+int(l)]...)
buf = append(buf, branchData[pos:pos+int(l)]...)
pos += int(l)
}
}

bitset ^= bit
}

return newData, nil
return bytes.Clone(buf), buf, nil
}

// IsComplete determines whether given branch data is complete, meaning that all information about all the children is present
Expand Down Expand Up @@ -1286,8 +1289,8 @@ func (m *BranchMerger) Merge(branch1 BranchData, branch2 BranchData) (BranchData
}
pos1 += n
if len(branch1) < pos1+int(l) {
fmt.Printf("b1: %x %v\n", branch1, branch1)
fmt.Printf("b2: %x\n", branch2)
//fmt.Printf("b1: %x %v\n", branch1, branch1)
//fmt.Printf("b2: %x\n", branch2)
return nil, fmt.Errorf("MergeHexBranches branch1 is too small: expected at least %d got %d bytes", pos1+int(l), len(branch1))
}
if l > 0 {
Expand Down
Loading
Loading