Skip to content

Commit 25b4e57

Browse files
qdm12qdm12
committed
Do not add inlined nodes to proof nodes
1 parent 53e4a64 commit 25b4e57

File tree

2 files changed

+128
-30
lines changed

2 files changed

+128
-30
lines changed

‎lib/trie/proof/generate.go

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,8 @@ func Generate(rootHash []byte, fullKeys [][]byte, database Database) (
3939
hashesSeen := make(map[string]struct{})
4040
for _, fullKey := range fullKeys {
4141
fullKeyNibbles := codec.KeyLEToNibbles(fullKey)
42-
newEncodedProofNodes, err := walk(rootNode, fullKeyNibbles)
42+
const isRoot = true
43+
newEncodedProofNodes, err := walk(rootNode, fullKeyNibbles, isRoot)
4344
if err != nil {
4445
// Note we wrap the full key context here since walk is recursive and
4546
// may not be aware of the initial full key.
@@ -66,7 +67,7 @@ func Generate(rootHash []byte, fullKeys [][]byte, database Database) (
6667
return encodedProofNodes, nil
6768
}
6869

69-
func walk(parent *node.Node, fullKey []byte) (
70+
func walk(parent *node.Node, fullKey []byte, isRoot bool) (
7071
encodedProofNodes [][]byte, err error) {
7172
if parent == nil {
7273
if len(fullKey) == 0 {
@@ -82,7 +83,15 @@ func walk(parent *node.Node, fullKey []byte) (
8283
if err != nil {
8384
return nil, fmt.Errorf("encode node: %w", err)
8485
}
85-
encodedProofNodes = append(encodedProofNodes, encodingBuffer.Bytes())
86+
87+
if isRoot || encodingBuffer.Len() >= 32 {
88+
// Only add the root node encoding (whatever its length)
89+
// and child node encodings greater or equal to 32 bytes.
90+
// This is because child node encodings of less than 32 bytes
91+
// are inlined in the parent node encoding, so there is no need
92+
// to duplicate them in the proof generated.
93+
encodedProofNodes = append(encodedProofNodes, encodingBuffer.Bytes())
94+
}
8695

8796
nodeFound := len(fullKey) == 0 || bytes.Equal(parent.Key, fullKey)
8897
if nodeFound {
@@ -102,7 +111,8 @@ func walk(parent *node.Node, fullKey []byte) (
102111
childIndex := fullKey[commonLength]
103112
nextChild := parent.Children[childIndex]
104113
nextFullKey := fullKey[commonLength+1:]
105-
deeperEncodedProofNodes, err := walk(nextChild, nextFullKey)
114+
isRoot = false
115+
deeperEncodedProofNodes, err := walk(nextChild, nextFullKey, isRoot)
106116
if err != nil {
107117
return nil, err // note: do not wrap since this is recursive
108118
}

‎lib/trie/proof/generate_test.go

Lines changed: 114 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,9 @@ func Test_Generate(t *testing.T) {
2525
someHash[i] = byte(i)
2626
}
2727

28+
largeValue := generateBytes(t, 40)
29+
assertLongEncoding(t, node.Node{Value: largeValue})
30+
2831
testCases := map[string]struct {
2932
rootHash []byte
3033
fullKeysNibbles [][]byte
@@ -124,19 +127,26 @@ func Test_Generate(t *testing.T) {
124127
},
125128
databaseBuilder: func(ctrl *gomock.Controller) Database {
126129
mockDatabase := NewMockDatabase(ctrl)
127-
encodedRoot := encodeNode(t, node.Node{
130+
131+
rootNode := node.Node{
128132
Key: []byte{1, 2},
129133
Value: []byte{2},
130134
Children: padRightChildren([]*node.Node{
131135
nil, nil, nil,
132136
{ // full key 1, 2, 3, 4
133137
Key: []byte{4},
134-
Value: []byte{4},
138+
Value: largeValue,
135139
},
136140
}),
137-
})
141+
}
142+
138143
mockDatabase.EXPECT().Get(someHash).
139-
Return(encodedRoot, nil)
144+
Return(encodeNode(t, rootNode), nil)
145+
146+
encodedChild := encodeNode(t, *rootNode.Children[3])
147+
mockDatabase.EXPECT().Get(blake2b(t, encodedChild)).
148+
Return(encodedChild, nil)
149+
140150
return mockDatabase
141151
},
142152
encodedProofNodes: [][]byte{
@@ -147,13 +157,13 @@ func Test_Generate(t *testing.T) {
147157
nil, nil, nil,
148158
{
149159
Key: []byte{4},
150-
Value: []byte{4},
160+
Value: largeValue,
151161
},
152162
}),
153163
}),
154164
encodeNode(t, node.Node{
155165
Key: []byte{4},
156-
Value: []byte{4},
166+
Value: largeValue,
157167
}),
158168
},
159169
},
@@ -162,31 +172,42 @@ func Test_Generate(t *testing.T) {
162172
fullKeysNibbles: [][]byte{
163173
{1, 2, 3, 4},
164174
{1, 2, 4, 4},
165-
{1, 2, 5, 4},
175+
{1, 2, 5, 5},
166176
},
167177
databaseBuilder: func(ctrl *gomock.Controller) Database {
168178
mockDatabase := NewMockDatabase(ctrl)
169-
encodedRoot := encodeNode(t, node.Node{
179+
180+
rootNode := node.Node{
170181
Key: []byte{1, 2},
171182
Value: []byte{2},
172183
Children: padRightChildren([]*node.Node{
173184
nil, nil, nil,
174185
{ // full key 1, 2, 3, 4
175186
Key: []byte{4},
176-
Value: []byte{4},
187+
Value: largeValue,
177188
},
178189
{ // full key 1, 2, 4, 4
179190
Key: []byte{4},
180-
Value: []byte{4},
191+
Value: largeValue,
181192
},
182-
{ // full key 1, 2, 5, 4
183-
Key: []byte{4},
184-
Value: []byte{5},
193+
{ // full key 1, 2, 5, 5
194+
Key: []byte{5},
195+
Value: largeValue,
185196
},
186197
}),
187-
})
198+
}
199+
188200
mockDatabase.EXPECT().Get(someHash).
189-
Return(encodedRoot, nil)
201+
Return(encodeNode(t, rootNode), nil)
202+
203+
encodedLargeChild1 := encodeNode(t, *rootNode.Children[3])
204+
mockDatabase.EXPECT().Get(blake2b(t, encodedLargeChild1)).
205+
Return(encodedLargeChild1, nil).Times(2)
206+
207+
encodedLargeChild2 := encodeNode(t, *rootNode.Children[5])
208+
mockDatabase.EXPECT().Get(blake2b(t, encodedLargeChild2)).
209+
Return(encodedLargeChild2, nil)
210+
190211
return mockDatabase
191212
},
192213
encodedProofNodes: [][]byte{
@@ -197,25 +218,25 @@ func Test_Generate(t *testing.T) {
197218
nil, nil, nil,
198219
{ // full key 1, 2, 3, 4
199220
Key: []byte{4},
200-
Value: []byte{4},
221+
Value: largeValue,
201222
},
202223
{ // full key 1, 2, 4, 4
203224
Key: []byte{4},
204-
Value: []byte{4},
225+
Value: largeValue,
205226
},
206-
{ // full key 1, 2, 5, 4
207-
Key: []byte{4},
208-
Value: []byte{5},
227+
{ // full key 1, 2, 5, 5
228+
Key: []byte{5},
229+
Value: largeValue,
209230
},
210231
}),
211232
}),
212233
encodeNode(t, node.Node{
213234
Key: []byte{4},
214-
Value: []byte{4},
235+
Value: largeValue,
215236
}),
216237
encodeNode(t, node.Node{
217-
Key: []byte{4},
218-
Value: []byte{5},
238+
Key: []byte{5},
239+
Value: largeValue,
219240
}),
220241
},
221242
},
@@ -248,9 +269,13 @@ func Test_Generate(t *testing.T) {
248269
func Test_walk(t *testing.T) {
249270
t.Parallel()
250271

272+
largeValue := generateBytes(t, 40)
273+
assertLongEncoding(t, node.Node{Value: largeValue})
274+
251275
testCases := map[string]struct {
252276
parent *node.Node
253277
fullKey []byte // nibbles
278+
isRoot bool
254279
encodedProofNodes [][]byte
255280
errWrapped error
256281
errMessage string
@@ -268,6 +293,7 @@ func Test_walk(t *testing.T) {
268293
Key: []byte{1, 2},
269294
Value: []byte{1},
270295
},
296+
isRoot: true,
271297
encodedProofNodes: [][]byte{encodeNode(t, node.Node{
272298
Key: []byte{1, 2},
273299
Value: []byte{1},
@@ -311,6 +337,7 @@ func Test_walk(t *testing.T) {
311337
},
312338
}),
313339
},
340+
isRoot: true,
314341
encodedProofNodes: [][]byte{
315342
encodeNode(t, node.Node{
316343
Key: []byte{1, 2},
@@ -366,6 +393,7 @@ func Test_walk(t *testing.T) {
366393
}),
367394
},
368395
fullKey: []byte{1, 2},
396+
isRoot: true,
369397
encodedProofNodes: [][]byte{
370398
encodeNode(t, node.Node{
371399
Key: []byte{1, 2},
@@ -379,6 +407,64 @@ func Test_walk(t *testing.T) {
379407
}),
380408
},
381409
},
410+
"branch and matching search key for small leaf encoding": {
411+
parent: &node.Node{
412+
Key: []byte{1, 2},
413+
Value: []byte{3},
414+
Children: padRightChildren([]*node.Node{
415+
{ // full key 1, 2, 0, 1, 2
416+
Key: []byte{1, 2},
417+
Value: []byte{3},
418+
},
419+
}),
420+
},
421+
fullKey: []byte{1, 2, 0, 1, 2},
422+
isRoot: true,
423+
encodedProofNodes: [][]byte{
424+
encodeNode(t, node.Node{
425+
Key: []byte{1, 2},
426+
Value: []byte{3},
427+
Children: padRightChildren([]*node.Node{
428+
{ // full key 1, 2, 0, 1, 2
429+
Key: []byte{1, 2},
430+
Value: []byte{3},
431+
},
432+
}),
433+
}),
434+
// Note the leaf encoding is not added since its encoding
435+
// is less than 32 bytes.
436+
},
437+
},
438+
"branch and matching search key for large leaf encoding": {
439+
parent: &node.Node{
440+
Key: []byte{1, 2},
441+
Value: []byte{3},
442+
Children: padRightChildren([]*node.Node{
443+
{ // full key 1, 2, 0, 1, 2
444+
Key: []byte{1, 2},
445+
Value: largeValue,
446+
},
447+
}),
448+
},
449+
fullKey: []byte{1, 2, 0, 1, 2},
450+
isRoot: true,
451+
encodedProofNodes: [][]byte{
452+
encodeNode(t, node.Node{
453+
Key: []byte{1, 2},
454+
Value: []byte{3},
455+
Children: padRightChildren([]*node.Node{
456+
{ // full key 1, 2, 0, 1, 2
457+
Key: []byte{1, 2},
458+
Value: largeValue,
459+
},
460+
}),
461+
}),
462+
encodeNode(t, node.Node{
463+
Key: []byte{1, 2},
464+
Value: largeValue,
465+
}),
466+
},
467+
},
382468
"key not found at deeper level": {
383469
parent: &node.Node{
384470
Key: []byte{1, 2},
@@ -406,6 +492,7 @@ func Test_walk(t *testing.T) {
406492
}),
407493
},
408494
fullKey: []byte{1, 2, 0x04},
495+
isRoot: true,
409496
encodedProofNodes: [][]byte{
410497
encodeNode(t, node.Node{
411498
Key: []byte{1, 2},
@@ -426,7 +513,7 @@ func Test_walk(t *testing.T) {
426513
t.Run(name, func(t *testing.T) {
427514
t.Parallel()
428515

429-
encodedProofNodes, err := walk(testCase.parent, testCase.fullKey)
516+
encodedProofNodes, err := walk(testCase.parent, testCase.fullKey, testCase.isRoot)
430517

431518
assert.ErrorIs(t, err, testCase.errWrapped)
432519
if testCase.errWrapped != nil {
@@ -514,12 +601,13 @@ func Benchmark_walk(b *testing.B) {
514601
longestKeyNibbles := codec.KeyLEToNibbles(longestKeyLE)
515602

516603
rootNode := trie.RootNode()
517-
encodedProofNodes, err := walk(rootNode, longestKeyNibbles)
604+
const isRoot = true
605+
encodedProofNodes, err := walk(rootNode, longestKeyNibbles, isRoot)
518606
require.NoError(b, err)
519607
require.Equal(b, len(encodedProofNodes), trieDepth)
520608

521609
b.ResetTimer()
522610
for i := 0; i < b.N; i++ {
523-
_, _ = walk(rootNode, longestKeyNibbles)
611+
_, _ = walk(rootNode, longestKeyNibbles, isRoot)
524612
}
525613
}

0 commit comments

Comments
 (0)