Skip to content

Commit

Permalink
feat(trie): use scale encoder (#2930)
Browse files Browse the repository at this point in the history
  • Loading branch information
@qdm12
qdm12 authored Nov 14, 2022
1 parent 9093db2 commit e3dc108
Show file tree
Hide file tree
Showing 5 changed files with 67 additions and 147 deletions.
30 changes: 8 additions & 22 deletions internal/trie/node/branch_encode.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -105,43 +105,29 @@ func encodeChildrenSequentially(children []*Node, buffer io.Writer) (err error)
for i, child := range children {
err = encodeChild(child, buffer)
if err != nil {
return fmt.Errorf("cannot encode child at index %d: %w", i, err)
return fmt.Errorf("encoding child at index %d: %w", i, err)
}
}
return nil
}

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

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

_, err = buffer.Write(scaleEncodedChildHash)
encoder := scale.NewEncoder(buffer)
err = encoder.Encode(merkleValue)
if err != nil {
return fmt.Errorf("failed to write child to buffer: %w", err)
return fmt.Errorf("scale encoding Merkle value: %w", err)
}

return nil
}

// scaleEncodeHash hashes the node (blake2b sum on encoded value)
// and then SCALE encodes it. This is used to encode children
// nodes of branches.
func scaleEncodeHash(node *Node) (encoding []byte, err error) {
_, merkleValue, err := node.EncodeAndHash()
if err != nil {
return nil, fmt.Errorf("encoding and hashing %s: %w", node.Kind(), err)
}

encoding, err = scale.Marshal(merkleValue)
if err != nil {
return nil, fmt.Errorf("cannot scale encode hashed %s: %w", node.Kind(), err)
}

return encoding, nil
}
129 changes: 37 additions & 92 deletions internal/trie/node/branch_encode_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -118,23 +118,6 @@ func Test_encodeChildrenOpportunisticParallel(t *testing.T) {
errMessage: "cannot write encoding of child at index 11: " +
"test error",
},
"branch encoding": {
// Note this may run in parallel or not depending on other tests
// running in parallel.
children: []*Node{
{
Key: []byte{1},
Children: []*Node{
{Key: []byte{1}, SubValue: []byte{2}},
},
},
},
writes: []writeCall{
{
written: []byte{36, 129, 1, 1, 0, 16, 65, 1, 4, 2},
},
},
},
}

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

buffer := bytes.NewBuffer(nil)

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

require.NoError(t, err)
Expand Down Expand Up @@ -210,9 +195,8 @@ func Test_encodeChildrenSequentially(t *testing.T) {
{Key: []byte{1}, SubValue: []byte{2}},
},
writes: []writeCall{
{
written: []byte{16, 65, 1, 4, 2},
},
{written: []byte{16}},
{written: []byte{65, 1, 4, 2}},
},
},
"last child not nil": {
Expand All @@ -223,9 +207,8 @@ func Test_encodeChildrenSequentially(t *testing.T) {
{Key: []byte{1}, SubValue: []byte{2}},
},
writes: []writeCall{
{
written: []byte{16, 65, 1, 4, 2},
},
{written: []byte{16}},
{written: []byte{65, 1, 4, 2}},
},
},
"first two children not nil": {
Expand All @@ -234,12 +217,10 @@ func Test_encodeChildrenSequentially(t *testing.T) {
{Key: []byte{3}, SubValue: []byte{4}},
},
writes: []writeCall{
{
written: []byte{16, 65, 1, 4, 2},
},
{
written: []byte{16, 65, 3, 4, 4},
},
{written: []byte{16}},
{written: []byte{65, 1, 4, 2}},
{written: []byte{16}},
{written: []byte{65, 3, 4, 4}},
},
},
"encoding error": {
Expand All @@ -252,13 +233,13 @@ func Test_encodeChildrenSequentially(t *testing.T) {
},
writes: []writeCall{
{
written: []byte{16, 65, 1, 4, 2},
written: []byte{16},
err: errTest,
},
},
wrappedErr: errTest,
errMessage: "cannot encode child at index 11: " +
"failed to write child to buffer: test error",
errMessage: "encoding child at index 11: " +
"scale encoding Merkle value: test error",
},
}

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

testCases := map[string]struct {
child *Node
writeCall bool
write writeCall
writes []writeCall
wrappedErr error
errMessage string
}{
Expand All @@ -308,31 +288,30 @@ func Test_encodeChild(t *testing.T) {
child: &Node{
Children: make([]*Node, ChildrenCapacity),
},
writeCall: true,
write: writeCall{
written: []byte{12, 128, 0, 0},
writes: []writeCall{
{written: []byte{12}},
{written: []byte{128, 0, 0}},
},
},
"buffer write error": {
"scale encoding error": {
child: &Node{
Children: make([]*Node, ChildrenCapacity),
},
writeCall: true,
write: writeCall{
written: []byte{12, 128, 0, 0},
writes: []writeCall{{
written: []byte{12},
err: errTest,
},
}},
wrappedErr: errTest,
errMessage: "failed to write child to buffer: test error",
errMessage: "scale encoding Merkle value: test error",
},
"leaf child": {
child: &Node{
Key: []byte{1},
SubValue: []byte{2},
},
writeCall: true,
write: writeCall{
written: []byte{16, 65, 1, 4, 2},
writes: []writeCall{
{written: []byte{16}},
{written: []byte{65, 1, 4, 2}},
},
},
"branch child": {
Expand All @@ -345,9 +324,9 @@ func Test_encodeChild(t *testing.T) {
},
},
},
writeCall: true,
write: writeCall{
written: []byte{44, 193, 1, 4, 0, 4, 2, 16, 65, 5, 4, 6},
writes: []writeCall{
{written: []byte{44}},
{written: []byte{193, 1, 4, 0, 4, 2, 16, 65, 5, 4, 6}},
},
},
}
Expand All @@ -360,10 +339,15 @@ func Test_encodeChild(t *testing.T) {

buffer := NewMockWriter(ctrl)

if testCase.writeCall {
buffer.EXPECT().
Write(testCase.write.written).
Return(testCase.write.n, testCase.write.err)
var previousCall *gomock.Call
for _, write := range testCase.writes {
call := buffer.EXPECT().
Write(write.written).
Return(write.n, write.err)
if previousCall != nil {
call.After(previousCall)
}
previousCall = call
}

err := encodeChild(testCase.child, buffer)
Expand All @@ -377,42 +361,3 @@ func Test_encodeChild(t *testing.T) {
})
}
}

func Test_scaleEncodeHash(t *testing.T) {
t.Parallel()

testCases := map[string]struct {
node *Node
encoding []byte
wrappedErr error
errMessage string
}{
"branch": {
node: &Node{
Key: []byte{1, 2},
SubValue: []byte{3, 4},
Children: []*Node{
nil, nil, {Key: []byte{9}, SubValue: []byte{1}},
},
},
encoding: []byte{0x30, 0xc2, 0x12, 0x4, 0x0, 0x8, 0x3, 0x4, 0x10, 0x41, 0x9, 0x4, 0x1},
},
}

for name, testCase := range testCases {
testCase := testCase
t.Run(name, func(t *testing.T) {
t.Parallel()

encoding, err := scaleEncodeHash(testCase.node)

if testCase.wrappedErr != nil {
assert.ErrorIs(t, err, testCase.wrappedErr)
assert.EqualError(t, err, testCase.errMessage)
} else {
require.NoError(t, err)
}
assert.Equal(t, testCase.encoding, encoding)
})
}
}
10 changes: 3 additions & 7 deletions internal/trie/node/encode.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,14 +40,10 @@ func (n *Node) Encode(buffer Buffer) (err error) {
// Only encode node value if the node is a leaf or
// the node is a branch with a non empty value.
if !nodeIsBranch || (nodeIsBranch && n.SubValue != nil) {
encodedValue, err := scale.Marshal(n.SubValue) // TODO scale encoder to write to buffer
encoder := scale.NewEncoder(buffer)
err = encoder.Encode(n.SubValue)
if err != nil {
return fmt.Errorf("cannot scale encode value: %w", err)
}

_, err = buffer.Write(encodedValue)
if err != nil {
return fmt.Errorf("cannot write scale encoded value to buffer: %w", err)
return fmt.Errorf("scale encoding value: %w", err)
}
}

Expand Down
42 changes: 17 additions & 25 deletions internal/trie/node/encode_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -73,12 +73,12 @@ func Test_Node_Encode(t *testing.T) {
written: []byte{0x01, 0x23},
},
{
written: []byte{12, 4, 5, 6},
written: []byte{12},
err: errTest,
},
},
wrappedErr: errTest,
errMessage: "cannot write scale encoded value to buffer: test error",
errMessage: "scale encoding value: test error",
},
"leaf success": {
node: &Node{
Expand All @@ -89,12 +89,9 @@ func Test_Node_Encode(t *testing.T) {
{
written: []byte{leafVariant.bits | 3}, // partial key length 3
},
{
written: []byte{0x01, 0x23},
},
{
written: []byte{12, 4, 5, 6},
},
{written: []byte{0x01, 0x23}},
{written: []byte{12}},
{written: []byte{4, 5, 6}},
},
expectedEncoding: []byte{1, 2, 3},
},
Expand All @@ -103,15 +100,10 @@ func Test_Node_Encode(t *testing.T) {
Key: []byte{1, 2, 3},
},
writes: []writeCall{
{
written: []byte{leafVariant.bits | 3}, // partial key length 3
},
{
written: []byte{0x01, 0x23},
},
{
written: []byte{0},
},
{written: []byte{leafVariant.bits | 3}}, // partial key length 3
{written: []byte{0x01, 0x23}}, // partial key
{written: []byte{0}}, // node value encoded length
{written: nil}, // node value
},
expectedEncoding: []byte{1, 2, 3},
},
Expand Down Expand Up @@ -191,12 +183,12 @@ func Test_Node_Encode(t *testing.T) {
written: []byte{136, 0},
},
{ // value
written: []byte{4, 100},
written: []byte{4},
err: errTest,
},
},
wrappedErr: errTest,
errMessage: "cannot write scale encoded value to buffer: test error",
errMessage: "scale encoding value: test error",
},
"buffer write error for children encoding": {
node: &Node{
Expand All @@ -217,9 +209,9 @@ func Test_Node_Encode(t *testing.T) {
{ // children bitmap
written: []byte{136, 0},
},
{ // value
written: []byte{4, 100},
},
// value
{written: []byte{4}},
{written: []byte{100}},
{ // children
written: []byte{16, 65, 9, 4, 1},
err: errTest,
Expand Down Expand Up @@ -249,9 +241,9 @@ func Test_Node_Encode(t *testing.T) {
{ // children bitmap
written: []byte{136, 0},
},
{ // value
written: []byte{4, 100},
},
// value
{written: []byte{4}},
{written: []byte{100}},
{ // first children
written: []byte{16, 65, 9, 4, 1},
},
Expand Down
3 changes: 2 additions & 1 deletion lib/trie/trie_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -343,7 +343,8 @@ func Test_encodeRoot(t *testing.T) {
writeCalls: []writeCall{
{written: []byte{66}},
{written: []byte{18}},
{written: []byte{4, 1}},
{written: []byte{4}},
{written: []byte{1}},
},
expectedRoot: &Node{
Key: []byte{1, 2},
Expand Down

0 comments on commit e3dc108

Please sign in to comment.