Skip to content

Commit

Permalink
Merge leaf and branch header encoding
Browse files Browse the repository at this point in the history
  • Loading branch information
@qdm12
qdm12 committed May 3, 2022
1 parent c2af35d commit 257277e
Show file tree
Hide file tree
Showing 2 changed files with 51 additions and 125 deletions.
64 changes: 14 additions & 50 deletions internal/trie/node/header.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
112 changes: 37 additions & 75 deletions internal/trie/node/header_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,50 +10,50 @@ 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),
},
writes: []writeCall{
{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),
},
writes: []writeCall{
{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),
},
writes: []writeCall{
{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),
},
Expand All @@ -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),
},
Expand All @@ -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),
},
Expand All @@ -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{
Expand All @@ -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),
},
Expand All @@ -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{
Expand All @@ -175,34 +137,34 @@ 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{
{written: []byte{0x7f}},
{written: []byte{0x0}},
},
},
"key of length 64": {
leaf: &Node{
"leaf with key of length 64": {
node: &Node{
Key: make([]byte, 64),
},
writes: []writeCall{
{written: []byte{0x7f}},
{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{
Expand All @@ -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{
Expand Down Expand Up @@ -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 {
Expand Down

0 comments on commit 257277e

Please sign in to comment.