From 257277e9cfe27bb575a7f2d8a4fccecf756fdf85 Mon Sep 17 00:00:00 2001 From: Quentin McGaw Date: Tue, 3 May 2022 17:46:56 +0000 Subject: [PATCH] Merge leaf and branch header encoding --- internal/trie/node/header.go | 64 ++++------------- internal/trie/node/header_test.go | 112 ++++++++++-------------------- 2 files changed, 51 insertions(+), 125 deletions(-) diff --git a/internal/trie/node/header.go b/internal/trie/node/header.go index adf51e2c080..05d268a2d5e 100644 --- a/internal/trie/node/header.go +++ b/internal/trie/node/header.go @@ -15,68 +15,32 @@ const ( nodeHeaderShift = 6 ) +// encodeHeader writes the encoded header for the node. func encodeHeader(node *Node, writer io.Writer) (err error) { - switch node.Type() { - case Leaf: - return encodeLeafHeader(node, writer) - case Branch: - return encodeBranchHeader(node, writer) - default: - panic("header encoding not implemented") - } -} - -// encodeBranchHeader writes the encoded header for the branch. -func encodeBranchHeader(branch *Node, writer io.Writer) (err error) { var header byte - if branch.Value == nil { - header = branchHeaderByte << nodeHeaderShift - } else { - header = branchWithValueHeaderByte << nodeHeaderShift - } - - if len(branch.Key) >= keyLenOffset { - header = header | keyLenOffset - _, err = writer.Write([]byte{header}) - if err != nil { - return err - } - - err = encodeKeyLength(len(branch.Key), writer) - if err != nil { - return err - } - } else { - header = header | byte(len(branch.Key)) - _, err = writer.Write([]byte{header}) - if err != nil { - return err + if node.Type() == Leaf { + header = leafHeaderByte + } else { // branch + if node.Value == nil { + header = branchHeaderByte + } else { + header = branchWithValueHeaderByte } } + header <<= nodeHeaderShift - return nil -} - -// encodeLeafHeader writes the encoded header for the leaf. -func encodeLeafHeader(leaf *Node, writer io.Writer) (err error) { - header := leafHeaderByte << nodeHeaderShift - - if len(leaf.Key) < 63 { - header |= byte(len(leaf.Key)) + if len(node.Key) < keyLenOffset { + header |= byte(len(node.Key)) _, err = writer.Write([]byte{header}) return err } - header |= keyLenOffset + header = header | keyLenOffset _, err = writer.Write([]byte{header}) if err != nil { return err } - err = encodeKeyLength(len(leaf.Key), writer) - if err != nil { - return err - } - - return nil + err = encodeKeyLength(len(node.Key), writer) + return err } diff --git a/internal/trie/node/header_test.go b/internal/trie/node/header_test.go index 16147d75de8..1ed826483af 100644 --- a/internal/trie/node/header_test.go +++ b/internal/trie/node/header_test.go @@ -10,23 +10,23 @@ import ( "github.com/stretchr/testify/assert" ) -func Test_encodeBranchHeader(t *testing.T) { +func Test_encodeHeader(t *testing.T) { testCases := map[string]struct { - branch *Node + node *Node writes []writeCall errWrapped error errMessage string }{ - "no key": { - branch: &Node{ + "branch with no key": { + node: &Node{ Children: make([]*Node, ChildrenCapacity), }, writes: []writeCall{ {written: []byte{0x80}}, }, }, - "with value": { - branch: &Node{ + "branch with value": { + node: &Node{ Value: []byte{}, Children: make([]*Node, ChildrenCapacity), }, @@ -34,8 +34,8 @@ func Test_encodeBranchHeader(t *testing.T) { {written: []byte{0xc0}}, }, }, - "key of length 30": { - branch: &Node{ + "branch with key of length 30": { + node: &Node{ Key: make([]byte, 30), Children: make([]*Node, ChildrenCapacity), }, @@ -43,8 +43,8 @@ func Test_encodeBranchHeader(t *testing.T) { {written: []byte{0x9e}}, }, }, - "key of length 62": { - branch: &Node{ + "branch with key of length 62": { + node: &Node{ Key: make([]byte, 62), Children: make([]*Node, ChildrenCapacity), }, @@ -52,8 +52,8 @@ func Test_encodeBranchHeader(t *testing.T) { {written: []byte{0xbe}}, }, }, - "key of length 63": { - branch: &Node{ + "branch with key of length 63": { + node: &Node{ Key: make([]byte, 63), Children: make([]*Node, ChildrenCapacity), }, @@ -62,8 +62,8 @@ func Test_encodeBranchHeader(t *testing.T) { {written: []byte{0x0}}, }, }, - "key of length 64": { - branch: &Node{ + "branch with key of length 64": { + node: &Node{ Key: make([]byte, 64), Children: make([]*Node, ChildrenCapacity), }, @@ -72,8 +72,8 @@ func Test_encodeBranchHeader(t *testing.T) { {written: []byte{0x1}}, }, }, - "key too big": { - branch: &Node{ + "branch with key too big": { + node: &Node{ Key: make([]byte, 65535+63), Children: make([]*Node, ChildrenCapacity), }, @@ -83,8 +83,8 @@ func Test_encodeBranchHeader(t *testing.T) { errWrapped: ErrPartialKeyTooBig, errMessage: "partial key length cannot be larger than or equal to 2^16: 65535", }, - "small key length write error": { - branch: &Node{ + "branch with small key length write error": { + node: &Node{ Children: make([]*Node, ChildrenCapacity), }, writes: []writeCall{ @@ -96,8 +96,8 @@ func Test_encodeBranchHeader(t *testing.T) { errWrapped: errTest, errMessage: "test error", }, - "long key length write error": { - branch: &Node{ + "branch with long key length write error": { + node: &Node{ Key: make([]byte, 64), Children: make([]*Node, ChildrenCapacity), }, @@ -110,60 +110,22 @@ func Test_encodeBranchHeader(t *testing.T) { errWrapped: errTest, errMessage: "test error", }, - } - - for name, testCase := range testCases { - testCase := testCase - t.Run(name, func(t *testing.T) { - t.Parallel() - ctrl := gomock.NewController(t) - - writer := NewMockWriter(ctrl) - var previousCall *gomock.Call - for _, write := range testCase.writes { - call := writer.EXPECT(). - Write(write.written). - Return(write.n, write.err) - - if previousCall != nil { - call.After(previousCall) - } - previousCall = call - } - - err := encodeBranchHeader(testCase.branch, writer) - - assert.ErrorIs(t, err, testCase.errWrapped) - if testCase.errWrapped != nil { - assert.EqualError(t, err, testCase.errMessage) - } - }) - } -} - -func Test_encodeLeafHeader(t *testing.T) { - testCases := map[string]struct { - leaf *Node - writes []writeCall - errWrapped error - errMessage string - }{ - "no key": { - leaf: &Node{}, + "leaf with no key": { + node: &Node{}, writes: []writeCall{ {written: []byte{0x40}}, }, }, - "key of length 30": { - leaf: &Node{ + "leaf with key of length 30": { + node: &Node{ Key: make([]byte, 30), }, writes: []writeCall{ {written: []byte{0x5e}}, }, }, - "short key write error": { - leaf: &Node{ + "leaf with short key write error": { + node: &Node{ Key: make([]byte, 30), }, writes: []writeCall{ @@ -175,16 +137,16 @@ func Test_encodeLeafHeader(t *testing.T) { errWrapped: errTest, errMessage: errTest.Error(), }, - "key of length 62": { - leaf: &Node{ + "leaf with key of length 62": { + node: &Node{ Key: make([]byte, 62), }, writes: []writeCall{ {written: []byte{0x7e}}, }, }, - "key of length 63": { - leaf: &Node{ + "leaf with key of length 63": { + node: &Node{ Key: make([]byte, 63), }, writes: []writeCall{ @@ -192,8 +154,8 @@ func Test_encodeLeafHeader(t *testing.T) { {written: []byte{0x0}}, }, }, - "key of length 64": { - leaf: &Node{ + "leaf with key of length 64": { + node: &Node{ Key: make([]byte, 64), }, writes: []writeCall{ @@ -201,8 +163,8 @@ func Test_encodeLeafHeader(t *testing.T) { {written: []byte{0x1}}, }, }, - "long key first byte write error": { - leaf: &Node{ + "leaf with long key first byte write error": { + node: &Node{ Key: make([]byte, 63), }, writes: []writeCall{ @@ -214,8 +176,8 @@ func Test_encodeLeafHeader(t *testing.T) { errWrapped: errTest, errMessage: errTest.Error(), }, - "key too big": { - leaf: &Node{ + "leaf with key too big": { + node: &Node{ Key: make([]byte, 65535+63), }, writes: []writeCall{ @@ -245,7 +207,7 @@ func Test_encodeLeafHeader(t *testing.T) { previousCall = call } - err := encodeLeafHeader(testCase.leaf, writer) + err := encodeHeader(testCase.node, writer) assert.ErrorIs(t, err, testCase.errWrapped) if testCase.errWrapped != nil {