Skip to content

Commit e3dc108

Browse files
qdm12qdm12
authored
feat(trie): use scale encoder (#2930)
1 parent 9093db2 commit e3dc108

File tree

5 files changed

+67
-147
lines changed

5 files changed

+67
-147
lines changed

‎internal/trie/node/branch_encode.go

Lines changed: 8 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -105,43 +105,29 @@ func encodeChildrenSequentially(children []*Node, buffer io.Writer) (err error)
105105
for i, child := range children {
106106
err = encodeChild(child, buffer)
107107
if err != nil {
108-
return fmt.Errorf("cannot encode child at index %d: %w", i, err)
108+
return fmt.Errorf("encoding child at index %d: %w", i, err)
109109
}
110110
}
111111
return nil
112112
}
113113

114+
// encodeChild computes the Merkle value of the node
115+
// and then SCALE encodes it to the given buffer.
114116
func encodeChild(child *Node, buffer io.Writer) (err error) {
115117
if child == nil {
116118
return nil
117119
}
118120

119-
scaleEncodedChildHash, err := scaleEncodeHash(child)
121+
_, merkleValue, err := child.EncodeAndHash()
120122
if err != nil {
121-
return fmt.Errorf("failed to hash and scale encode child: %w", err)
123+
return fmt.Errorf("computing %s Merkle value: %w", child.Kind(), err)
122124
}
123125

124-
_, err = buffer.Write(scaleEncodedChildHash)
126+
encoder := scale.NewEncoder(buffer)
127+
err = encoder.Encode(merkleValue)
125128
if err != nil {
126-
return fmt.Errorf("failed to write child to buffer: %w", err)
129+
return fmt.Errorf("scale encoding Merkle value: %w", err)
127130
}
128131

129132
return nil
130133
}
131-
132-
// scaleEncodeHash hashes the node (blake2b sum on encoded value)
133-
// and then SCALE encodes it. This is used to encode children
134-
// nodes of branches.
135-
func scaleEncodeHash(node *Node) (encoding []byte, err error) {
136-
_, merkleValue, err := node.EncodeAndHash()
137-
if err != nil {
138-
return nil, fmt.Errorf("encoding and hashing %s: %w", node.Kind(), err)
139-
}
140-
141-
encoding, err = scale.Marshal(merkleValue)
142-
if err != nil {
143-
return nil, fmt.Errorf("cannot scale encode hashed %s: %w", node.Kind(), err)
144-
}
145-
146-
return encoding, nil
147-
}

‎internal/trie/node/branch_encode_test.go

Lines changed: 37 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -118,23 +118,6 @@ func Test_encodeChildrenOpportunisticParallel(t *testing.T) {
118118
errMessage: "cannot write encoding of child at index 11: " +
119119
"test error",
120120
},
121-
"branch encoding": {
122-
// Note this may run in parallel or not depending on other tests
123-
// running in parallel.
124-
children: []*Node{
125-
{
126-
Key: []byte{1},
127-
Children: []*Node{
128-
{Key: []byte{1}, SubValue: []byte{2}},
129-
},
130-
},
131-
},
132-
writes: []writeCall{
133-
{
134-
written: []byte{36, 129, 1, 1, 0, 16, 65, 1, 4, 2},
135-
},
136-
},
137-
},
138121
}
139122

140123
for name, testCase := range testCases {
@@ -179,6 +162,8 @@ func Test_encodeChildrenOpportunisticParallel(t *testing.T) {
179162

180163
buffer := bytes.NewBuffer(nil)
181164

165+
// Note this may run in parallel or not depending on other tests
166+
// running in parallel.
182167
err := encodeChildrenOpportunisticParallel(children, buffer)
183168

184169
require.NoError(t, err)
@@ -210,9 +195,8 @@ func Test_encodeChildrenSequentially(t *testing.T) {
210195
{Key: []byte{1}, SubValue: []byte{2}},
211196
},
212197
writes: []writeCall{
213-
{
214-
written: []byte{16, 65, 1, 4, 2},
215-
},
198+
{written: []byte{16}},
199+
{written: []byte{65, 1, 4, 2}},
216200
},
217201
},
218202
"last child not nil": {
@@ -223,9 +207,8 @@ func Test_encodeChildrenSequentially(t *testing.T) {
223207
{Key: []byte{1}, SubValue: []byte{2}},
224208
},
225209
writes: []writeCall{
226-
{
227-
written: []byte{16, 65, 1, 4, 2},
228-
},
210+
{written: []byte{16}},
211+
{written: []byte{65, 1, 4, 2}},
229212
},
230213
},
231214
"first two children not nil": {
@@ -234,12 +217,10 @@ func Test_encodeChildrenSequentially(t *testing.T) {
234217
{Key: []byte{3}, SubValue: []byte{4}},
235218
},
236219
writes: []writeCall{
237-
{
238-
written: []byte{16, 65, 1, 4, 2},
239-
},
240-
{
241-
written: []byte{16, 65, 3, 4, 4},
242-
},
220+
{written: []byte{16}},
221+
{written: []byte{65, 1, 4, 2}},
222+
{written: []byte{16}},
223+
{written: []byte{65, 3, 4, 4}},
243224
},
244225
},
245226
"encoding error": {
@@ -252,13 +233,13 @@ func Test_encodeChildrenSequentially(t *testing.T) {
252233
},
253234
writes: []writeCall{
254235
{
255-
written: []byte{16, 65, 1, 4, 2},
236+
written: []byte{16},
256237
err: errTest,
257238
},
258239
},
259240
wrappedErr: errTest,
260-
errMessage: "cannot encode child at index 11: " +
261-
"failed to write child to buffer: test error",
241+
errMessage: "encoding child at index 11: " +
242+
"scale encoding Merkle value: test error",
262243
},
263244
}
264245

@@ -298,8 +279,7 @@ func Test_encodeChild(t *testing.T) {
298279

299280
testCases := map[string]struct {
300281
child *Node
301-
writeCall bool
302-
write writeCall
282+
writes []writeCall
303283
wrappedErr error
304284
errMessage string
305285
}{
@@ -308,31 +288,30 @@ func Test_encodeChild(t *testing.T) {
308288
child: &Node{
309289
Children: make([]*Node, ChildrenCapacity),
310290
},
311-
writeCall: true,
312-
write: writeCall{
313-
written: []byte{12, 128, 0, 0},
291+
writes: []writeCall{
292+
{written: []byte{12}},
293+
{written: []byte{128, 0, 0}},
314294
},
315295
},
316-
"buffer write error": {
296+
"scale encoding error": {
317297
child: &Node{
318298
Children: make([]*Node, ChildrenCapacity),
319299
},
320-
writeCall: true,
321-
write: writeCall{
322-
written: []byte{12, 128, 0, 0},
300+
writes: []writeCall{{
301+
written: []byte{12},
323302
err: errTest,
324-
},
303+
}},
325304
wrappedErr: errTest,
326-
errMessage: "failed to write child to buffer: test error",
305+
errMessage: "scale encoding Merkle value: test error",
327306
},
328307
"leaf child": {
329308
child: &Node{
330309
Key: []byte{1},
331310
SubValue: []byte{2},
332311
},
333-
writeCall: true,
334-
write: writeCall{
335-
written: []byte{16, 65, 1, 4, 2},
312+
writes: []writeCall{
313+
{written: []byte{16}},
314+
{written: []byte{65, 1, 4, 2}},
336315
},
337316
},
338317
"branch child": {
@@ -345,9 +324,9 @@ func Test_encodeChild(t *testing.T) {
345324
},
346325
},
347326
},
348-
writeCall: true,
349-
write: writeCall{
350-
written: []byte{44, 193, 1, 4, 0, 4, 2, 16, 65, 5, 4, 6},
327+
writes: []writeCall{
328+
{written: []byte{44}},
329+
{written: []byte{193, 1, 4, 0, 4, 2, 16, 65, 5, 4, 6}},
351330
},
352331
},
353332
}
@@ -360,10 +339,15 @@ func Test_encodeChild(t *testing.T) {
360339

361340
buffer := NewMockWriter(ctrl)
362341

363-
if testCase.writeCall {
364-
buffer.EXPECT().
365-
Write(testCase.write.written).
366-
Return(testCase.write.n, testCase.write.err)
342+
var previousCall *gomock.Call
343+
for _, write := range testCase.writes {
344+
call := buffer.EXPECT().
345+
Write(write.written).
346+
Return(write.n, write.err)
347+
if previousCall != nil {
348+
call.After(previousCall)
349+
}
350+
previousCall = call
367351
}
368352

369353
err := encodeChild(testCase.child, buffer)
@@ -377,42 +361,3 @@ func Test_encodeChild(t *testing.T) {
377361
})
378362
}
379363
}
380-
381-
func Test_scaleEncodeHash(t *testing.T) {
382-
t.Parallel()
383-
384-
testCases := map[string]struct {
385-
node *Node
386-
encoding []byte
387-
wrappedErr error
388-
errMessage string
389-
}{
390-
"branch": {
391-
node: &Node{
392-
Key: []byte{1, 2},
393-
SubValue: []byte{3, 4},
394-
Children: []*Node{
395-
nil, nil, {Key: []byte{9}, SubValue: []byte{1}},
396-
},
397-
},
398-
encoding: []byte{0x30, 0xc2, 0x12, 0x4, 0x0, 0x8, 0x3, 0x4, 0x10, 0x41, 0x9, 0x4, 0x1},
399-
},
400-
}
401-
402-
for name, testCase := range testCases {
403-
testCase := testCase
404-
t.Run(name, func(t *testing.T) {
405-
t.Parallel()
406-
407-
encoding, err := scaleEncodeHash(testCase.node)
408-
409-
if testCase.wrappedErr != nil {
410-
assert.ErrorIs(t, err, testCase.wrappedErr)
411-
assert.EqualError(t, err, testCase.errMessage)
412-
} else {
413-
require.NoError(t, err)
414-
}
415-
assert.Equal(t, testCase.encoding, encoding)
416-
})
417-
}
418-
}

‎internal/trie/node/encode.go

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -40,14 +40,10 @@ func (n *Node) Encode(buffer Buffer) (err error) {
4040
// Only encode node value if the node is a leaf or
4141
// the node is a branch with a non empty value.
4242
if !nodeIsBranch || (nodeIsBranch && n.SubValue != nil) {
43-
encodedValue, err := scale.Marshal(n.SubValue) // TODO scale encoder to write to buffer
43+
encoder := scale.NewEncoder(buffer)
44+
err = encoder.Encode(n.SubValue)
4445
if err != nil {
45-
return fmt.Errorf("cannot scale encode value: %w", err)
46-
}
47-
48-
_, err = buffer.Write(encodedValue)
49-
if err != nil {
50-
return fmt.Errorf("cannot write scale encoded value to buffer: %w", err)
46+
return fmt.Errorf("scale encoding value: %w", err)
5147
}
5248
}
5349

‎internal/trie/node/encode_test.go

Lines changed: 17 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -73,12 +73,12 @@ func Test_Node_Encode(t *testing.T) {
7373
written: []byte{0x01, 0x23},
7474
},
7575
{
76-
written: []byte{12, 4, 5, 6},
76+
written: []byte{12},
7777
err: errTest,
7878
},
7979
},
8080
wrappedErr: errTest,
81-
errMessage: "cannot write scale encoded value to buffer: test error",
81+
errMessage: "scale encoding value: test error",
8282
},
8383
"leaf success": {
8484
node: &Node{
@@ -89,12 +89,9 @@ func Test_Node_Encode(t *testing.T) {
8989
{
9090
written: []byte{leafVariant.bits | 3}, // partial key length 3
9191
},
92-
{
93-
written: []byte{0x01, 0x23},
94-
},
95-
{
96-
written: []byte{12, 4, 5, 6},
97-
},
92+
{written: []byte{0x01, 0x23}},
93+
{written: []byte{12}},
94+
{written: []byte{4, 5, 6}},
9895
},
9996
expectedEncoding: []byte{1, 2, 3},
10097
},
@@ -103,15 +100,10 @@ func Test_Node_Encode(t *testing.T) {
103100
Key: []byte{1, 2, 3},
104101
},
105102
writes: []writeCall{
106-
{
107-
written: []byte{leafVariant.bits | 3}, // partial key length 3
108-
},
109-
{
110-
written: []byte{0x01, 0x23},
111-
},
112-
{
113-
written: []byte{0},
114-
},
103+
{written: []byte{leafVariant.bits | 3}}, // partial key length 3
104+
{written: []byte{0x01, 0x23}}, // partial key
105+
{written: []byte{0}}, // node value encoded length
106+
{written: nil}, // node value
115107
},
116108
expectedEncoding: []byte{1, 2, 3},
117109
},
@@ -191,12 +183,12 @@ func Test_Node_Encode(t *testing.T) {
191183
written: []byte{136, 0},
192184
},
193185
{ // value
194-
written: []byte{4, 100},
186+
written: []byte{4},
195187
err: errTest,
196188
},
197189
},
198190
wrappedErr: errTest,
199-
errMessage: "cannot write scale encoded value to buffer: test error",
191+
errMessage: "scale encoding value: test error",
200192
},
201193
"buffer write error for children encoding": {
202194
node: &Node{
@@ -217,9 +209,9 @@ func Test_Node_Encode(t *testing.T) {
217209
{ // children bitmap
218210
written: []byte{136, 0},
219211
},
220-
{ // value
221-
written: []byte{4, 100},
222-
},
212+
// value
213+
{written: []byte{4}},
214+
{written: []byte{100}},
223215
{ // children
224216
written: []byte{16, 65, 9, 4, 1},
225217
err: errTest,
@@ -249,9 +241,9 @@ func Test_Node_Encode(t *testing.T) {
249241
{ // children bitmap
250242
written: []byte{136, 0},
251243
},
252-
{ // value
253-
written: []byte{4, 100},
254-
},
244+
// value
245+
{written: []byte{4}},
246+
{written: []byte{100}},
255247
{ // first children
256248
written: []byte{16, 65, 9, 4, 1},
257249
},

‎lib/trie/trie_test.go

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -343,7 +343,8 @@ func Test_encodeRoot(t *testing.T) {
343343
writeCalls: []writeCall{
344344
{written: []byte{66}},
345345
{written: []byte{18}},
346-
{written: []byte{4, 1}},
346+
{written: []byte{4}},
347+
{written: []byte{1}},
347348
},
348349
expectedRoot: &Node{
349350
Key: []byte{1, 2},

0 commit comments

Comments
 (0)