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 56069ef
Show file tree
Hide file tree
Showing 3 changed files with 63 additions and 134 deletions.
8 changes: 4 additions & 4 deletions internal/trie/node/decode.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,13 +36,13 @@ func Decode(reader io.Reader) (n *Node, err error) {

nodeTypeHeaderByte := header >> 6
switch nodeTypeHeaderByte {
case leafHeaderByte:
case leafHeader:
n, err = decodeLeaf(reader, header)
if err != nil {
return nil, fmt.Errorf("cannot decode leaf: %w", err)
}
return n, nil
case branchHeaderByte, branchWithValueHeaderByte:
case branchHeader, branchWithValueHeader:
n, err = decodeBranch(reader, header)
if err != nil {
return nil, fmt.Errorf("cannot decode branch: %w", err)
Expand Down Expand Up @@ -78,7 +78,7 @@ func decodeBranch(reader io.Reader, header byte) (node *Node, err error) {
sd := scale.NewDecoder(reader)

nodeType := header >> 6
if nodeType == branchWithValueHeaderByte {
if nodeType == branchWithValueHeader {
var value []byte
// branch w/ value
err := sd.Decode(&value)
Expand All @@ -103,7 +103,7 @@ func decodeBranch(reader io.Reader, header byte) (node *Node, err error) {
// Handle inlined leaf nodes.
const hashLength = 32
nodeTypeHeaderByte := hash[0] >> 6
if nodeTypeHeaderByte == leafHeaderByte && len(hash) < hashLength {
if nodeTypeHeaderByte == leafHeader && len(hash) < hashLength {
leaf, err := decodeLeaf(bytes.NewReader(hash[1:]), hash[0])
if err != nil {
return nil, fmt.Errorf("%w: at index %d: %s",
Expand Down
77 changes: 22 additions & 55 deletions internal/trie/node/header.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,75 +8,42 @@ import (
)

const (
leafHeaderByte byte = 0x1
branchHeaderByte byte = 2
branchWithValueHeaderByte byte = 3
keyLenOffset = 0x3f
nodeHeaderShift = 6
leafHeader byte = 1 // 01
branchHeader byte = 2 // 10
branchWithValueHeader byte = 3 // 11
)

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")
}
}
const (
keyLenOffset = 0x3f
nodeHeaderShift = 6
)

// encodeBranchHeader writes the encoded header for the branch.
func encodeBranchHeader(branch *Node, writer io.Writer) (err error) {
// encodeHeader writes the encoded header for the node.
func encodeHeader(node *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 = leafHeader
} else { // branch
if node.Value == nil {
header = branchHeader
} else {
header = branchWithValueHeader
}
}
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 56069ef

Please sign in to comment.