feat(trie): Add trie v1 new headers support (#3295)

dot/state/mocks_generate_test.go

@@ -7,3 +7,4 @@ package state
 //go:generate mockgen -destination=mocks_runtime_test.go -package $GOPACKAGE github.com/ChainSafe/gossamer/lib/runtime Instance
 //go:generate mockgen -destination=mock_gauge_test.go -package $GOPACKAGE github.com/prometheus/client_golang/prometheus Gauge
 //go:generate mockgen -destination=mock_counter_test.go -package $GOPACKAGE github.com/prometheus/client_golang/prometheus Counter
+//go:generate mockgen -destination=db_getter_mocks_test.go -package=$GOPACKAGE github.com/ChainSafe/gossamer/lib/trie DBGetter

dot/state/storage_test.go

@@ -178,13 +178,15 @@ func TestGetStorageChildAndGetStorageFromChild(t *testing.T) {
 		genHeader.Hash(),
 		"0",
 	))
+	dbGetter := NewMockDBGetter(ctrl)
+	dbGetter.EXPECT().Get(gomock.Any()).Times(0)
 
 	trieRoot := &node.Node{
 		PartialKey:   []byte{1, 2},
 		StorageValue: []byte{3, 4},
 		Dirty:        true,
 	}
-	testChildTrie := trie.NewTrie(trieRoot)
+	testChildTrie := trie.NewTrie(trieRoot, dbGetter)
 
 	testChildTrie.Put([]byte("keyInsidechild"), []byte("voila"))
 

dot/state/tries_test.go

@@ -48,8 +48,11 @@ func Test_Tries_SetEmptyTrie(t *testing.T) {
 
 func Test_Tries_SetTrie(t *testing.T) {
 	t.Parallel()
+	ctrl := gomock.NewController(t)
+	dbGetter := NewMockDBGetter(ctrl)
+	dbGetter.EXPECT().Get(gomock.Any()).Times(0)
 
-	tr := trie.NewTrie(&node.Node{PartialKey: []byte{1}})
+	tr := trie.NewTrie(&node.Node{PartialKey: []byte{1}}, dbGetter)
 
 	tries := NewTries()
 	tries.SetTrie(tr)
@@ -188,6 +191,9 @@ func Test_Tries_delete(t *testing.T) {
 }
 func Test_Tries_get(t *testing.T) {
 	t.Parallel()
+	ctrl := gomock.NewController(t)
+	dbGetter := NewMockDBGetter(ctrl)
+	dbGetter.EXPECT().Get(gomock.Any()).Times(0)
 
 	testCases := map[string]struct {
 		tries *Tries
@@ -200,14 +206,14 @@ func Test_Tries_get(t *testing.T) {
 					{1, 2, 3}: trie.NewTrie(&node.Node{
 						PartialKey:   []byte{1, 2, 3},
 						StorageValue: []byte{1},
-					}),
+					}, dbGetter),
 				},
 			},
 			root: common.Hash{1, 2, 3},
 			trie: trie.NewTrie(&node.Node{
 				PartialKey:   []byte{1, 2, 3},
 				StorageValue: []byte{1},
-			}),
+			}, dbGetter),
 		},
 		"not_found_in_map": {
 			// similar to not found in database

internal/trie/node/copy.go

@@ -90,6 +90,7 @@ func (n *Node) Copy(settings CopySettings) *Node {
 	if settings.CopyStorageValue && n.StorageValue != nil {
 		cpy.StorageValue = make([]byte, len(n.StorageValue))
 		copy(cpy.StorageValue, n.StorageValue)
+		cpy.HashedValue = n.HashedValue
 	}
 
 	if settings.CopyMerkleValue {

internal/trie/node/copy_test.go

@@ -110,6 +110,37 @@ func Test_Node_Copy(t *testing.T) {
 				MerkleValue: []byte{5},
 			},
 		},
+		"deep_copy_branch_with_hashed_values": {
+			node: &Node{
+				PartialKey:   []byte{1, 2},
+				StorageValue: []byte{3, 4},
+				HashedValue:  true,
+				Children: padRightChildren([]*Node{
+					nil, nil, {
+						PartialKey:   []byte{9},
+						StorageValue: []byte{1},
+						HashedValue:  true,
+					},
+				}),
+				Dirty:       true,
+				MerkleValue: []byte{5},
+			},
+			settings: DeepCopySettings,
+			expectedNode: &Node{
+				PartialKey:   []byte{1, 2},
+				StorageValue: []byte{3, 4},
+				HashedValue:  true,
+				Children: padRightChildren([]*Node{
+					nil, nil, {
+						PartialKey:   []byte{9},
+						StorageValue: []byte{1},
+						HashedValue:  true,
+					},
+				}),
+				Dirty:       true,
+				MerkleValue: []byte{5},
+			},
+		},
 		"non_empty_leaf": {
 			node: &Node{
 				PartialKey:   []byte{1, 2},
@@ -139,6 +170,23 @@ func Test_Node_Copy(t *testing.T) {
 				MerkleValue:  []byte{5},
 			},
 		},
+		"deep_copy_leaf_with_hashed_value": {
+			node: &Node{
+				PartialKey:   []byte{1, 2},
+				StorageValue: []byte{3, 4},
+				HashedValue:  true,
+				Dirty:        true,
+				MerkleValue:  []byte{5},
+			},
+			settings: DeepCopySettings,
+			expectedNode: &Node{
+				PartialKey:   []byte{1, 2},
+				StorageValue: []byte{3, 4},
+				HashedValue:  true,
+				Dirty:        true,
+				MerkleValue:  []byte{5},
+			},
+		},
 	}
 
 	for name, testCase := range testCases {

internal/trie/node/decode.go

@@ -9,23 +9,24 @@ import (
 	"fmt"
 	"io"
 
+	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/pkg/scale"
 )
 
 var (
 	// ErrDecodeStorageValue is defined since no sentinel error is defined
 	// in the scale package.
-	// TODO remove once the following issue is done:
-	// /~https://github.com/ChainSafe/gossamer/issues/2631 .
-	ErrDecodeStorageValue = errors.New("cannot decode storage value")
-	ErrReadChildrenBitmap = errors.New("cannot read children bitmap")
+	ErrDecodeStorageValue        = errors.New("cannot decode storage value")
+	ErrDecodeHashedStorageValue  = errors.New("cannot decode hashed storage value")
+	ErrDecodeHashedValueTooShort = errors.New("hashed storage value too short")
+	ErrReadChildrenBitmap        = errors.New("cannot read children bitmap")
 	// ErrDecodeChildHash is defined since no sentinel error is defined
 	// in the scale package.
-	// TODO remove once the following issue is done:
-	// /~https://github.com/ChainSafe/gossamer/issues/2631 .
 	ErrDecodeChildHash = errors.New("cannot decode child hash")
 )
 
+const hashLength = common.HashLength
+
 // Decode decodes a node from a reader.
 // The encoding format is documented in the README.md
 // of this package, and specified in the Polkadot spec at
@@ -39,21 +40,22 @@ func Decode(reader io.Reader) (n *Node, err error) {
 	}
 
 	switch variant {
-	case leafVariant.bits:
-		n, err = decodeLeaf(reader, partialKeyLength)
+	case emptyVariant:
+		return nil, nil //nolint:nilnil
+	case leafVariant, leafWithHashedValueVariant:
+		n, err = decodeLeaf(reader, variant, partialKeyLength)
 		if err != nil {
 			return nil, fmt.Errorf("cannot decode leaf: %w", err)
 		}
 		return n, nil
-	case branchVariant.bits, branchWithValueVariant.bits:
+	case branchVariant, branchWithValueVariant, branchWithHashedValueVariant:
 		n, err = decodeBranch(reader, variant, partialKeyLength)
 		if err != nil {
 			return nil, fmt.Errorf("cannot decode branch: %w", err)
 		}
 		return n, nil
 	default:
-		// this is a programming error, an unknown node variant
-		// should be caught by decodeHeader.
+		// this is a programming error, an unknown node variant should be caught by decodeHeader.
 		panic(fmt.Sprintf("not implemented for node variant %08b", variant))
 	}
 }
@@ -63,7 +65,7 @@ func Decode(reader io.Reader) (n *Node, err error) {
 // reconstructing the child nodes from the encoding. This function instead stubs where the
 // children are known to be with an empty leaf. The children nodes hashes are then used to
 // find other storage values using the persistent database.
-func decodeBranch(reader io.Reader, variant byte, partialKeyLength uint16) (
+func decodeBranch(reader io.Reader, variant variant, partialKeyLength uint16) (
 	node *Node, err error) {
 	node = &Node{
 		Children: make([]*Node, ChildrenCapacity),
@@ -82,11 +84,21 @@ func decodeBranch(reader io.Reader, variant byte, partialKeyLength uint16) (
 
 	sd := scale.NewDecoder(reader)
 
-	if variant == branchWithValueVariant.bits {
+	switch variant {
+	case branchWithValueVariant:
 		err := sd.Decode(&node.StorageValue)
 		if err != nil {
 			return nil, fmt.Errorf("%w: %s", ErrDecodeStorageValue, err)
 		}
+	case branchWithHashedValueVariant:
+		hashedValue, err := decodeHashedValue(reader)
+		if err != nil {
+			return nil, err
+		}
+		node.StorageValue = hashedValue
+		node.HashedValue = true
+	default:
+		// Ignored
 	}
 
 	for i := 0; i < ChildrenCapacity; i++ {
@@ -101,7 +113,6 @@ func decodeBranch(reader io.Reader, variant byte, partialKeyLength uint16) (
 				ErrDecodeChildHash, i, err)
 		}
 
-		const hashLength = 32
 		childNode := &Node{
 			MerkleValue: hash,
 		}
@@ -123,7 +134,7 @@ func decodeBranch(reader io.Reader, variant byte, partialKeyLength uint16) (
 }
 
 // decodeLeaf reads from a reader and decodes to a leaf node.
-func decodeLeaf(reader io.Reader, partialKeyLength uint16) (node *Node, err error) {
+func decodeLeaf(reader io.Reader, variant variant, partialKeyLength uint16) (node *Node, err error) {
 	node = &Node{}
 
 	node.PartialKey, err = decodeKey(reader, partialKeyLength)
@@ -132,10 +143,34 @@ func decodeLeaf(reader io.Reader, partialKeyLength uint16) (node *Node, err erro
 	}
 
 	sd := scale.NewDecoder(reader)
+
+	if variant == leafWithHashedValueVariant {
+		hashedValue, err := decodeHashedValue(reader)
+		if err != nil {
+			return nil, err
+		}
+		node.StorageValue = hashedValue
+		node.HashedValue = true
+		return node, nil
+	}
+
 	err = sd.Decode(&node.StorageValue)
 	if err != nil {
 		return nil, fmt.Errorf("%w: %s", ErrDecodeStorageValue, err)
 	}
 
 	return node, nil
 }
+
+func decodeHashedValue(reader io.Reader) ([]byte, error) {
+	buffer := make([]byte, hashLength)
+	n, err := reader.Read(buffer)
+	if err != nil {
+		return nil, fmt.Errorf("%w: %s", ErrDecodeStorageValue, err)
+	}
+	if n < hashLength {
+		return nil, fmt.Errorf("%w: expected %d, got: %d", ErrDecodeHashedValueTooShort, hashLength, n)
+	}
+
+	return buffer, nil
+}