Merge leaf and branch header encoding

Author: qdm12

internal/trie/node/decode.go

@@ -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)
@@ -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)
@@ -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",

internal/trie/node/header.go

@@ -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
 }

internal/trie/node/header_test.go

@@ -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),
 			},
@@ -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,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{
@@ -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 {