Skip to content

Commit 7dd6944

Browse files
AskAlexSharovsudeepdino008
authored andcommitted
replaceShortenedKeysInBranch: less re-allocs (#19584)
old `aux` buffer often did re-alloc (16children * 32hashlen = 368 > 256) ``` ┌───────────────────┬────────┬────────┬───────────┐ │ Pattern │ ns/op │ B/op │ allocs/op │ ├───────────────────┼────────┼────────┼───────────┤ │ fresh-make (old) │ 315 ns │ 1408 B │ 2 allocs │ ├───────────────────┼────────┼────────┼───────────┤ │ reuse-clone (new) │ 248 ns │ 704 B │ 1 alloc │ └───────────────────┴────────┴────────┴───────────┘ ```
1 parent d7120c3 commit 7dd6944

File tree

6 files changed

+101
-102
lines changed

6 files changed

+101
-102
lines changed

db/integrity/commitment_integirty.go

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -399,7 +399,7 @@ func checkDerefBranch(
399399
branchData := commitment.BranchData(branchValue)
400400
var integrityErr error
401401

402-
newBranchData, _ = branchData.ReplacePlainKeys(newBranchValueBuf[:0], func(key []byte, isStorage bool) ([]byte, error) {
402+
newBranchData, newBranchValueBuf, _ = branchData.ReplacePlainKeys(newBranchValueBuf[:0], func(key []byte, isStorage bool) ([]byte, error) {
403403
if trace {
404404
logger.Trace(
405405
"checking commitment deref for branch",
@@ -1023,6 +1023,7 @@ func checkStateCorrespondenceBase(ctx context.Context, file state.VisibleFile, s
10231023
branchValueBuf := make([]byte, 0, datasize.MB.Bytes())
10241024
var branchKeys uint64
10251025
var integrityErr error
1026+
var branchDerefBuf []byte
10261027

10271028
for i := 0; commReader.HasNext(); i++ {
10281029
if i%1024 == 0 {
@@ -1104,7 +1105,8 @@ func checkStateCorrespondenceBase(ctx context.Context, file state.VisibleFile, s
11041105
}
11051106
// The callback returns nil (keep original key in output) because the result of ReplacePlainKeys
11061107
// is discarded (_): all side-effects happen before the return.
1107-
_, err := branchData.ReplacePlainKeys(nil, func(key []byte, isStorage bool) ([]byte, error) {
1108+
var err error
1109+
_, branchDerefBuf, err = branchData.ReplacePlainKeys(branchDerefBuf[:0], func(key []byte, isStorage bool) ([]byte, error) {
11081110
var checkErr error
11091111
if isStorage {
11101112
checkErr = checkKey(key, stoCollector, storageReader, length.Addr+length.Hash, "storage")
@@ -1239,6 +1241,7 @@ func checkStateCorrespondenceReverse(ctx context.Context, file state.VisibleFile
12391241
var branchKeys uint64
12401242
var integrityErr error
12411243
var extractedAccKeys, extractedStoKeys, skippedAccKeys, skippedStoKeys uint64
1244+
var branchDerefBuf []byte
12421245

12431246
// Phase 1: Walk commitment branches, collect extracted plain keys into ETL collectors.
12441247
for i := 0; commReader.HasNext(); i++ {
@@ -1282,7 +1285,8 @@ func checkStateCorrespondenceReverse(ctx context.Context, file state.VisibleFile
12821285
continue
12831286
}
12841287

1285-
_, err := branchData.ReplacePlainKeys(nil, func(key []byte, isStorage bool) ([]byte, error) {
1288+
var err error
1289+
_, branchDerefBuf, err = branchData.ReplacePlainKeys(branchDerefBuf[:0], func(key []byte, isStorage bool) ([]byte, error) {
12861290
if isStorage {
12871291
plainKey := key
12881292
if len(key) != length.Addr+length.Hash {
@@ -1621,6 +1625,7 @@ func extractCommitmentRefsToCollectors(ctx context.Context, file state.VisibleFi
16211625
branchKeyBuf := make([]byte, 0, 128)
16221626
branchValueBuf := make([]byte, 0, datasize.MB.Bytes())
16231627
plainKeyBuf := make([]byte, 0, length.Addr+length.Hash)
1628+
var branchDerefBuf []byte
16241629

16251630
for commReader.HasNext() {
16261631
select {
@@ -1644,7 +1649,8 @@ func extractCommitmentRefsToCollectors(ctx context.Context, file state.VisibleFi
16441649
continue
16451650
}
16461651

1647-
if _, err := branchData.ReplacePlainKeys(nil, func(key []byte, isStorage bool) ([]byte, error) { //nolint:gocritic
1652+
var err error
1653+
_, branchDerefBuf, err = branchData.ReplacePlainKeys(branchDerefBuf[:0], func(key []byte, isStorage bool) ([]byte, error) { //nolint:gocritic
16481654
if isStorage {
16491655
plainKey := key
16501656
if len(key) != length.Addr+length.Hash {
@@ -1677,7 +1683,8 @@ func extractCommitmentRefsToCollectors(ctx context.Context, file state.VisibleFi
16771683
}
16781684
}
16791685
return plainKey, nil
1680-
}); err != nil {
1686+
})
1687+
if err != nil {
16811688
return err
16821689
}
16831690
}
@@ -1765,6 +1772,7 @@ func checkHashVerification(ctx context.Context, file state.VisibleFile, stepSize
17651772

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

17691777
for item := range workCh {
17701778
select {
@@ -1782,7 +1790,9 @@ func checkHashVerification(ctx context.Context, file state.VisibleFile, stepSize
17821790

17831791
// We need branch data with plain keys for VerifyBranchHashes.
17841792
// Walk the branch to extract + resolve all keys and read values.
1785-
resolvedBranchData, err := branchData.ReplacePlainKeys(nil, func(key []byte, isStorage bool) ([]byte, error) {
1793+
var resolvedBranchData commitment.BranchData
1794+
var err error
1795+
resolvedBranchData, branchDerefBuf, err = branchData.ReplacePlainKeys(branchDerefBuf[:0], func(key []byte, isStorage bool) ([]byte, error) {
17861796
if isStorage {
17871797
plainKey := key
17881798
isRef := len(key) != length.Addr+length.Hash

db/state/aggregator.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1874,6 +1874,8 @@ type AggregatorRoTx struct {
18741874
d [kv.DomainLen]*DomainRoTx
18751875
iis []*InvertedIndexRoTx
18761876

1877+
branchDerefBuf []byte // scratch buffer reused across replaceShortenedKeysInBranch calls
1878+
18771879
_leakID uint64 // set only if TRACE_AGG=true
18781880
}
18791881

db/state/domain_committed.go

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -85,8 +85,8 @@ func (at *AggregatorRoTx) replaceShortenedKeysInBranch(prefix []byte, branch com
8585
}
8686
}
8787

88-
aux := make([]byte, 0, 256)
89-
return branch.ReplacePlainKeys(aux, func(key []byte, isStorage bool) ([]byte, error) {
88+
var result commitment.BranchData
89+
result, at.branchDerefBuf, err = branch.ReplacePlainKeys(at.branchDerefBuf[:0], func(key []byte, isStorage bool) ([]byte, error) {
9090
if isStorage {
9191
if len(key) == length.Addr+length.Hash {
9292
return nil, nil // save storage key as is
@@ -121,6 +121,10 @@ func (at *AggregatorRoTx) replaceShortenedKeysInBranch(prefix []byte, branch com
121121
}
122122
return apkBuf, nil
123123
})
124+
if err != nil {
125+
return nil, err
126+
}
127+
return result, nil
124128
}
125129

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

430-
temp, err := commitment.BranchData(valBuf).ReplacePlainKeys(dt.comBuf[:0], replacer)
434+
var branchData commitment.BranchData
435+
branchData, dt.comBuf, err = commitment.BranchData(valBuf).ReplacePlainKeys(dt.comBuf[:0], replacer)
431436
if err != nil {
432437
return nil, err
433438
}
434-
dt.comBuf = append(dt.comBuf[:0], temp...) // cover branch data case
435-
return dt.comBuf, nil
439+
return branchData, nil
436440
}
437441

438442
return vt, nil

execution/commitment/commitment.go

Lines changed: 41 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -871,137 +871,140 @@ func (branchData BranchData) String() string {
871871
return sb.String()
872872
}
873873

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

880883
var numBuf [binary.MaxVarintLen64]byte
881884
touchMap := binary.BigEndian.Uint16(branchData[0:])
882885
afterMap := binary.BigEndian.Uint16(branchData[2:])
883886
if touchMap&afterMap == 0 {
884-
return branchData, nil
887+
return branchData, buf, nil
885888
}
886889
pos := 4
887-
newData = append(newData[:0], branchData[:4]...)
890+
buf = append(buf[:0], branchData[:4]...)
888891
for bitset, j := touchMap&afterMap, 0; bitset != 0; j++ {
889892
bit := bitset & -bitset
890893
fields := cellFields(branchData[pos])
891-
newData = append(newData, byte(fields))
894+
buf = append(buf, byte(fields))
892895
pos++
893896
if fields&fieldExtension != 0 {
894897
l, n := binary.Uvarint(branchData[pos:])
895898
if n == 0 {
896-
return nil, errors.New("replacePlainKeys buffer too small for hashedKey len")
899+
return nil, buf, errors.New("replacePlainKeys buffer too small for hashedKey len")
897900
} else if n < 0 {
898-
return nil, errors.New("replacePlainKeys value overflow for hashedKey len")
901+
return nil, buf, errors.New("replacePlainKeys value overflow for hashedKey len")
899902
}
900-
newData = append(newData, branchData[pos:pos+n]...)
903+
buf = append(buf, branchData[pos:pos+n]...)
901904
pos += n
902905
if len(branchData) < pos+int(l) {
903-
return nil, fmt.Errorf("replacePlainKeys buffer too small for hashedKey: expected %d got %d", pos+int(l), len(branchData))
906+
return nil, buf, fmt.Errorf("replacePlainKeys buffer too small for hashedKey: expected %d got %d", pos+int(l), len(branchData))
904907
}
905908
if l > 0 {
906-
newData = append(newData, branchData[pos:pos+int(l)]...)
909+
buf = append(buf, branchData[pos:pos+int(l)]...)
907910
pos += int(l)
908911
}
909912
}
910913
if fields&fieldAccountAddr != 0 {
911914
l, n := binary.Uvarint(branchData[pos:])
912915
if n == 0 {
913-
return nil, errors.New("replacePlainKeys buffer too small for accountAddr len")
916+
return nil, buf, errors.New("replacePlainKeys buffer too small for accountAddr len")
914917
} else if n < 0 {
915-
return nil, errors.New("replacePlainKeys value overflow for accountAddr len")
918+
return nil, buf, errors.New("replacePlainKeys value overflow for accountAddr len")
916919
}
917920
pos += n
918921
if len(branchData) < pos+int(l) {
919-
return nil, fmt.Errorf("replacePlainKeys buffer too small for accountAddr: expected %d got %d", pos+int(l), len(branchData))
922+
return nil, buf, fmt.Errorf("replacePlainKeys buffer too small for accountAddr: expected %d got %d", pos+int(l), len(branchData))
920923
}
921924
if l > 0 {
922925
pos += int(l)
923926
}
924927
newKey, err := fn(branchData[pos-int(l):pos], false)
925928
if err != nil {
926-
return nil, err
929+
return nil, buf, err
927930
}
928931
if newKey == nil {
929932
// invariant: fn returns nil (keep original) only for plain addr keys (length.Addr bytes)
930-
newData = append(newData, branchData[pos-int(l)-n:pos]...)
933+
buf = append(buf, branchData[pos-int(l)-n:pos]...)
931934
} else {
932935
// invariant: newKey is a short reference (≤8 bytes) or full plain addr (length.Addr bytes)
933936
n = binary.PutUvarint(numBuf[:], uint64(len(newKey)))
934-
newData = append(newData, numBuf[:n]...)
935-
newData = append(newData, newKey...)
937+
buf = append(buf, numBuf[:n]...)
938+
buf = append(buf, newKey...)
936939
}
937940
}
938941
if fields&fieldStorageAddr != 0 {
939942
l, n := binary.Uvarint(branchData[pos:])
940943
if n == 0 {
941-
return nil, errors.New("replacePlainKeys buffer too small for storageAddr len")
944+
return nil, buf, errors.New("replacePlainKeys buffer too small for storageAddr len")
942945
} else if n < 0 {
943-
return nil, errors.New("replacePlainKeys value overflow for storageAddr len")
946+
return nil, buf, errors.New("replacePlainKeys value overflow for storageAddr len")
944947
}
945948
pos += n
946949
if len(branchData) < pos+int(l) {
947-
return nil, fmt.Errorf("replacePlainKeys buffer too small for storageAddr: expected %d got %d", pos+int(l), len(branchData))
950+
return nil, buf, fmt.Errorf("replacePlainKeys buffer too small for storageAddr: expected %d got %d", pos+int(l), len(branchData))
948951
}
949952
if l > 0 {
950953
pos += int(l)
951954
}
952955
newKey, err := fn(branchData[pos-int(l):pos], true)
953956
if err != nil {
954-
return nil, err
957+
return nil, buf, err
955958
}
956959
if newKey == nil {
957960
// invariant: fn returns nil (keep original) only for plain storage keys (length.Addr+length.Hash bytes)
958-
newData = append(newData, branchData[pos-int(l)-n:pos]...) // -n to include length
961+
buf = append(buf, branchData[pos-int(l)-n:pos]...) // -n to include length
959962
} else {
960963
// invariant: newKey is a short reference (≤8 bytes) or full plain storage key (length.Addr+length.Hash bytes)
961964
n = binary.PutUvarint(numBuf[:], uint64(len(newKey)))
962-
newData = append(newData, numBuf[:n]...)
963-
newData = append(newData, newKey...)
965+
buf = append(buf, numBuf[:n]...)
966+
buf = append(buf, newKey...)
964967
}
965968
}
966969
if fields&fieldHash != 0 {
967970
l, n := binary.Uvarint(branchData[pos:])
968971
if n == 0 {
969-
return nil, errors.New("replacePlainKeys buffer too small for hash len")
972+
return nil, buf, errors.New("replacePlainKeys buffer too small for hash len")
970973
} else if n < 0 {
971-
return nil, errors.New("replacePlainKeys value overflow for hash len")
974+
return nil, buf, errors.New("replacePlainKeys value overflow for hash len")
972975
}
973-
newData = append(newData, branchData[pos:pos+n]...)
976+
buf = append(buf, branchData[pos:pos+n]...)
974977
pos += n
975978
if len(branchData) < pos+int(l) {
976-
return nil, fmt.Errorf("replacePlainKeys buffer too small for hash: expected %d got %d", pos+int(l), len(branchData))
979+
return nil, buf, fmt.Errorf("replacePlainKeys buffer too small for hash: expected %d got %d", pos+int(l), len(branchData))
977980
}
978981
if l > 0 {
979-
newData = append(newData, branchData[pos:pos+int(l)]...)
982+
buf = append(buf, branchData[pos:pos+int(l)]...)
980983
pos += int(l)
981984
}
982985
}
983986
if fields&fieldStateHash != 0 {
984987
l, n := binary.Uvarint(branchData[pos:])
985988
if n == 0 {
986-
return nil, errors.New("replacePlainKeys buffer too small for acLeaf hash len")
989+
return nil, buf, errors.New("replacePlainKeys buffer too small for acLeaf hash len")
987990
} else if n < 0 {
988-
return nil, errors.New("replacePlainKeys value overflow for acLeafhash len")
991+
return nil, buf, errors.New("replacePlainKeys value overflow for acLeafhash len")
989992
}
990-
newData = append(newData, branchData[pos:pos+n]...)
993+
buf = append(buf, branchData[pos:pos+n]...)
991994
pos += n
992995
if len(branchData) < pos+int(l) {
993-
return nil, fmt.Errorf("replacePlainKeys buffer too small for LeafHash: expected %d got %d", pos+int(l), len(branchData))
996+
return nil, buf, fmt.Errorf("replacePlainKeys buffer too small for LeafHash: expected %d got %d", pos+int(l), len(branchData))
994997
}
995998
if l > 0 {
996-
newData = append(newData, branchData[pos:pos+int(l)]...)
999+
buf = append(buf, branchData[pos:pos+int(l)]...)
9971000
pos += int(l)
9981001
}
9991002
}
10001003

10011004
bitset ^= bit
10021005
}
10031006

1004-
return newData, nil
1007+
return bytes.Clone(buf), buf, nil
10051008
}
10061009

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

0 commit comments

Comments
 (0)