From e6878414d0faa3ddd2ccb8347b4cec3be0f6272c Mon Sep 17 00:00:00 2001 From: Quentin McGaw Date: Tue, 8 Nov 2022 16:09:23 +0000 Subject: [PATCH] feat(trie): use scale encoder --- internal/trie/node/branch_encode.go | 30 ++---- internal/trie/node/branch_encode_test.go | 129 +++++++---------------- internal/trie/node/encode.go | 10 +- internal/trie/node/encode_test.go | 42 +++----- lib/trie/trie_test.go | 3 +- 5 files changed, 67 insertions(+), 147 deletions(-) diff --git a/internal/trie/node/branch_encode.go b/internal/trie/node/branch_encode.go index 78d5550243..b9a45e0052 100644 --- a/internal/trie/node/branch_encode.go +++ b/internal/trie/node/branch_encode.go @@ -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 -} diff --git a/internal/trie/node/branch_encode_test.go b/internal/trie/node/branch_encode_test.go index 55386ff6f5..2b63c96aa2 100644 --- a/internal/trie/node/branch_encode_test.go +++ b/internal/trie/node/branch_encode_test.go @@ -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 { @@ -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) @@ -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": { @@ -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": { @@ -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": { @@ -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", }, } @@ -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 }{ @@ -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": { @@ -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}}, }, }, } @@ -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) @@ -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) - }) - } -} diff --git a/internal/trie/node/encode.go b/internal/trie/node/encode.go index b8b9f11b5f..16cd6dd930 100644 --- a/internal/trie/node/encode.go +++ b/internal/trie/node/encode.go @@ -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) } } diff --git a/internal/trie/node/encode_test.go b/internal/trie/node/encode_test.go index a81fb875de..9c20be4e69 100644 --- a/internal/trie/node/encode_test.go +++ b/internal/trie/node/encode_test.go @@ -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{ @@ -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}, }, @@ -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}, }, @@ -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{ @@ -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, @@ -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}, }, diff --git a/lib/trie/trie_test.go b/lib/trie/trie_test.go index d799fa9f66..99af820a2a 100644 --- a/lib/trie/trie_test.go +++ b/lib/trie/trie_test.go @@ -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},