fix(trie): decode inline child nodes (#2369)

Author: notbdu

internal/trie/node/decode.go

@@ -105,6 +105,18 @@ func decodeBranch(reader io.Reader, header byte) (branch *Branch, err error) {
 				ErrDecodeChildHash, i, err)
 		}
 
+		// Handle inlined leaf nodes.
+		const hashLength = 32
+		if Type(hash[0]>>6) == LeafType && len(hash) < hashLength {
+			leaf, err := decodeLeaf(bytes.NewReader(hash[1:]), hash[0])
+			if err != nil {
+				return nil, fmt.Errorf("%w: at index %d: %s",
+					ErrDecodeValue, i, err)
+			}
+			branch.Children[i] = leaf
+			continue
+		}
+
 		branch.Children[i] = &Leaf{
 			HashDigest: hash,
 		}

internal/trie/node/decode_test.go

@@ -95,6 +95,65 @@ func Test_Decode(t *testing.T) {
 				Dirty: true,
 			},
 		},
+		"branch with two inlined children": {
+			reader: bytes.NewReader(
+				[]byte{
+					158, // node type 2 (branch w/o value) and key length 30
+					// Key data start
+					195, 101, 195, 207, 89, 214,
+					113, 235, 114, 218, 14, 122,
+					65, 19, 196, 16, 2, 80, 95,
+					14, 123, 144, 18, 9, 107,
+					65, 196, 235, 58, 175,
+					// Key data end
+					148, 127, 110, 164, 41, 8, 0, 0, 104, 95, 15, 31, 5,
+					21, 244, 98, 205, 207, 132, 224, 241, 214, 4, 93, 252,
+					187, 32, 134, 92, 74, 43, 127, 1, 0, 0,
+				},
+			),
+			n: &Branch{
+				Key: []byte{
+					12, 3, 6, 5, 12, 3,
+					12, 15, 5, 9, 13, 6,
+					7, 1, 14, 11, 7, 2,
+					13, 10, 0, 14, 7, 10,
+					4, 1, 1, 3, 12, 4,
+				},
+				Children: [16]Node{
+					nil, nil, nil, nil,
+					&Leaf{
+						Key: []byte{
+							14, 7, 11, 9, 0, 1,
+							2, 0, 9, 6, 11, 4,
+							1, 12, 4, 14, 11,
+							3, 10, 10, 15, 9,
+							4, 7, 15, 6, 14,
+							10, 4, 2, 9,
+						},
+						Value: []byte{0, 0},
+						Dirty: true,
+					},
+					nil, nil, nil, nil,
+					&Leaf{
+						Key: []byte{
+							15, 1, 15, 0, 5, 1,
+							5, 15, 4, 6, 2, 12,
+							13, 12, 15, 8, 4,
+							14, 0, 15, 1, 13,
+							6, 0, 4, 5, 13,
+							15, 12, 11, 11,
+						},
+						Value: []byte{
+							134, 92, 74, 43,
+							127, 1, 0, 0,
+						},
+						Dirty: true,
+					},
+					nil, nil, nil, nil, nil, nil,
+				},
+				Dirty: true,
+			},
+		},
 	}
 
 	for name, testCase := range testCases {

internal/trie/node/encode_decode_test.go

@@ -43,7 +43,7 @@ func Test_Branch_Encode_Decode(t *testing.T) {
 				Dirty: true,
 			},
 		},
-		"branch with child": {
+		"branch with child leaf inline": {
 			branchToEncode: &Branch{
 				Key: []byte{5},
 				Children: [16]Node{
@@ -57,7 +57,50 @@ func Test_Branch_Encode_Decode(t *testing.T) {
 				Key: []byte{5},
 				Children: [16]Node{
 					&Leaf{
-						HashDigest: []byte{0x41, 0x9, 0x4, 0xa},
+						Key:   []byte{9},
+						Value: []byte{10},
+						Dirty: true,
+					},
+				},
+				Dirty: true,
+			},
+		},
+		"branch with child leaf hash": {
+			branchToEncode: &Branch{
+				Key: []byte{5},
+				Children: [16]Node{
+					&Leaf{
+						Key: []byte{
+							10, 11, 12, 13,
+							14, 15, 16, 17,
+							18, 19, 20, 21,
+							14, 15, 16, 17,
+							10, 11, 12, 13,
+							14, 15, 16, 17,
+						},
+						Value: []byte{
+							10, 11, 12, 13,
+							14, 15, 16, 17,
+							10, 11, 12, 13,
+							14, 15, 16, 17,
+							10, 11, 12, 13,
+						},
+					},
+				},
+			},
+			branchDecoded: &Branch{
+				Key: []byte{5},
+				Children: [16]Node{
+					&Leaf{
+						HashDigest: []byte{
+							2, 18, 48, 30, 98,
+							133, 244, 78, 70,
+							161, 196, 105, 228,
+							190, 159, 228, 199, 29,
+							254, 212, 160, 55, 199,
+							21, 186, 226, 204, 145,
+							132, 5, 39, 204,
+						},
 					},
 				},
 				Dirty: true,

lib/trie/database.go

@@ -186,6 +186,18 @@ func (t *Trie) load(db chaindb.Database, n Node) error {
 
 		hash := child.GetHash()
 
+		_, isLeaf := child.(*node.Leaf)
+		if len(hash) == 0 && isLeaf {
+			// node has already been loaded inline
+			// just set encoding + hash digest
+			_, _, err := child.EncodeAndHash(false)
+			if err != nil {
+				return err
+			}
+			child.SetDirty(false)
+			continue
+		}
+
 		encodedNode, err := db.Get(hash)
 		if err != nil {
 			return fmt.Errorf("cannot find child node key 0x%x in database: %w", hash, err)
@@ -330,12 +342,18 @@ func getFromDB(db chaindb.Database, n Node, key []byte) (
 
 	// childIndex is the nibble after the common prefix length in the key being searched.
 	childIndex := key[commonPrefixLength]
-	childWithHashOnly := branch.Children[childIndex]
-	if childWithHashOnly == nil {
+	child := branch.Children[childIndex]
+	if child == nil {
 		return nil, nil
 	}
 
-	childHash := childWithHashOnly.GetHash()
+	// Child can be either inlined or a hash pointer.
+	childHash := child.GetHash()
+	_, isLeaf := child.(*node.Leaf)
+	if len(childHash) == 0 && isLeaf {
+		return getFromDB(db, child, key[commonPrefixLength+1:])
+	}
+
 	encodedChild, err := db.Get(childHash)
 	if err != nil {
 		return nil, fmt.Errorf(