fix(lib/trie): Make sure writing and reading a trie to disk gives the same trie and cover more store/load child trie related test cases (#2302)

* test TestGetStorageChildAndGetStorageFromChild with non-empty trie
* test store and load of a trie that hash child tries
* tackle the case when encoding and hash is same, i.e., encoding is less than 32 bits
* don't put childTrie again inside the trie

Author: kishansagathiya

dot/state/storage.go

@@ -167,7 +167,7 @@ func (s *StorageState) loadTrie(root *common.Hash) (*trie.Trie, error) {
 
 	tr, err := s.LoadFromDB(*root)
 	if err != nil {
-		return nil, errTrieDoesNotExist(*root)
+		return nil, fmt.Errorf("trie does not exist at root %s: %w", *root, err)
 	}
 
 	return tr, nil

dot/state/storage_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/ChainSafe/gossamer/dot/state/pruner"
 	"github.com/ChainSafe/gossamer/dot/telemetry"
 	"github.com/ChainSafe/gossamer/dot/types"
+	"github.com/ChainSafe/gossamer/internal/trie/node"
 	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/lib/genesis"
 	runtime "github.com/ChainSafe/gossamer/lib/runtime/storage"
@@ -179,7 +180,8 @@ func TestGetStorageChildAndGetStorageFromChild(t *testing.T) {
 		"0",
 	))
 
-	testChildTrie := trie.NewEmptyTrie()
+	testChildTrie := trie.NewTrie(node.NewLeaf([]byte{1, 2}, []byte{3, 4}, true, 0))
+
 	testChildTrie.Put([]byte("keyInsidechild"), []byte("voila"))
 
 	err = genTrie.PutChild([]byte("keyToChild"), testChildTrie)

internal/trie/node/leaf.go

@@ -16,7 +16,7 @@ type Leaf struct {
 	// Partial key bytes in nibbles (0 to f in hexadecimal)
 	Key   []byte
 	Value []byte
-	// Dirty is true when the branch differs
+	// Dirty is true when the leaf differs
 	// from the node stored in the database.
 	Dirty      bool
 	HashDigest []byte

lib/trie/child_storage.go

@@ -16,16 +16,16 @@ var ChildStorageKeyPrefix = []byte(":child_storage:default:")
 var ErrChildTrieDoesNotExist = errors.New("child trie does not exist")
 
 // PutChild inserts a child trie into the main trie at key :child_storage:[keyToChild]
+// A child trie is added as a node (K, V) in the main trie. K is the child storage key
+// associated to the child trie, and V is the root hash of the child trie.
 func (t *Trie) PutChild(keyToChild []byte, child *Trie) error {
 	childHash, err := child.Hash()
 	if err != nil {
 		return err
 	}
-
 	key := append(ChildStorageKeyPrefix, keyToChild...)
-	value := [32]byte(childHash)
 
-	t.Put(key, value[:])
+	t.Put(key, childHash.ToBytes())
 	t.childTries[childHash] = child
 	return nil
 }
@@ -38,9 +38,7 @@ func (t *Trie) GetChild(keyToChild []byte) (*Trie, error) {
 		return nil, fmt.Errorf("%w at key 0x%x%x", ErrChildTrieDoesNotExist, ChildStorageKeyPrefix, keyToChild)
 	}
 
-	hash := [32]byte{}
-	copy(hash[:], childHash)
-	return t.childTries[common.Hash(hash)], nil
+	return t.childTries[common.BytesToHash(childHash)], nil
 }
 
 // PutIntoChild puts a key-value pair into the child trie located in the main trie at key :child_storage:[keyToChild]

lib/trie/database.go

@@ -150,8 +150,7 @@ func (t *Trie) Load(db chaindb.Database, rootHash common.Hash) error {
 		t.root = nil
 		return nil
 	}
-
-	rootHashBytes := rootHash[:]
+	rootHashBytes := rootHash.ToBytes()
 
 	encodedNode, err := db.Get(rootHashBytes)
 	if err != nil {
@@ -163,6 +162,7 @@ func (t *Trie) Load(db chaindb.Database, rootHash common.Hash) error {
 	if err != nil {
 		return fmt.Errorf("cannot decode root node: %w", err)
 	}
+
 	t.root = root
 	t.root.SetDirty(false)
 	t.root.SetEncodingAndHash(encodedNode, rootHashBytes)
@@ -185,6 +185,7 @@ func (t *Trie) load(db chaindb.Database, n Node) error {
 		}
 
 		hash := child.GetHash()
+
 		encodedNode, err := db.Get(hash)
 		if err != nil {
 			return fmt.Errorf("cannot find child node key 0x%x in database: %w", hash, err)
@@ -209,16 +210,17 @@ func (t *Trie) load(db chaindb.Database, n Node) error {
 	for _, key := range t.GetKeysWithPrefix(ChildStorageKeyPrefix) {
 		childTrie := NewEmptyTrie()
 		value := t.Get(key)
-		err := childTrie.Load(db, common.NewHash(value))
+		rootHash := common.BytesToHash(value)
+		err := childTrie.Load(db, rootHash)
 		if err != nil {
-			return fmt.Errorf("failed to load child trie with root hash=0x%x: %w", value, err)
+			return fmt.Errorf("failed to load child trie with root hash=%s: %w", rootHash, err)
 		}
 
-		err = t.PutChild(value, childTrie)
+		hash, err := childTrie.Hash()
 		if err != nil {
-			return fmt.Errorf("failed to insert child trie with root hash=0x%x into main trie: %w",
-				childTrie.root.GetHash(), err)
+			return fmt.Errorf("cannot hash chilld trie at key 0x%x: %w", key, err)
 		}
+		t.childTries[hash] = childTrie
 	}
 
 	return nil

lib/trie/database_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 
 	"github.com/ChainSafe/chaindb"
+	"github.com/ChainSafe/gossamer/internal/trie/node"
 	"github.com/ChainSafe/gossamer/lib/utils"
 	"github.com/stretchr/testify/require"
 )
@@ -21,7 +22,7 @@ func newTestDB(t *testing.T) chaindb.Database {
 }
 
 func TestTrie_DatabaseStoreAndLoad(t *testing.T) {
-	cases := [][]Test{
+	cases := [][]keyValues{
 		{
 			{key: []byte{0x01, 0x35}, value: []byte("pen")},
 			{key: []byte{0x01, 0x35, 0x79}, value: []byte("penguin")},
@@ -65,7 +66,7 @@ func TestTrie_DatabaseStoreAndLoad(t *testing.T) {
 		res := NewEmptyTrie()
 		err = res.Load(db, trie.MustHash())
 		require.NoError(t, err)
-		require.Equal(t, trie.MustHash(), res.MustHash())
+		require.Equal(t, trie.String(), res.String())
 
 		for _, test := range testCase {
 			val, err := GetFromDB(db, trie.MustHash(), test.key)
@@ -76,7 +77,7 @@ func TestTrie_DatabaseStoreAndLoad(t *testing.T) {
 }
 
 func TestTrie_WriteDirty_Put(t *testing.T) {
-	cases := [][]Test{
+	cases := [][]keyValues{
 		{
 			{key: []byte{0x01, 0x35}, value: []byte("pen")},
 			{key: []byte{0x01, 0x35, 0x79}, value: []byte("penguin")},
@@ -154,7 +155,7 @@ func TestTrie_WriteDirty_Put(t *testing.T) {
 }
 
 func TestTrie_WriteDirty_PutReplace(t *testing.T) {
-	cases := [][]Test{
+	cases := [][]keyValues{
 		{
 			{key: []byte{0x01, 0x35}, value: []byte("pen")},
 			{key: []byte{0x01, 0x35, 0x79}, value: []byte("penguin")},
@@ -217,7 +218,7 @@ func TestTrie_WriteDirty_PutReplace(t *testing.T) {
 }
 
 func TestTrie_WriteDirty_Delete(t *testing.T) {
-	cases := [][]Test{
+	cases := [][]keyValues{
 		{
 			{key: []byte{0x01, 0x35}, value: []byte("pen")},
 			{key: []byte{0x01, 0x35, 0x79}, value: []byte("penguin")},
@@ -284,7 +285,7 @@ func TestTrie_WriteDirty_Delete(t *testing.T) {
 }
 
 func TestTrie_WriteDirty_ClearPrefix(t *testing.T) {
-	cases := [][]Test{
+	cases := [][]keyValues{
 		{
 			{key: []byte{0x01, 0x35}, value: []byte("pen")},
 			{key: []byte{0x01, 0x35, 0x79}, value: []byte("penguin")},
@@ -337,7 +338,7 @@ func TestTrie_WriteDirty_ClearPrefix(t *testing.T) {
 }
 
 func TestTrie_GetFromDB(t *testing.T) {
-	cases := [][]Test{
+	cases := [][]keyValues{
 		{
 			{key: []byte{0x01, 0x35}, value: []byte("pen")},
 			{key: []byte{0x01, 0x35, 0x79}, value: []byte("penguin")},
@@ -387,3 +388,61 @@ func TestTrie_GetFromDB(t *testing.T) {
 		}
 	}
 }
+
+func TestStoreAndLoadWithChildTries(t *testing.T) {
+	keyValue := []keyValues{
+		{key: []byte{0xf2, 0x3}, value: []byte("f")},
+		{key: []byte{0x09, 0xd3}, value: []byte("noot")},
+		{key: []byte{0x07}, value: []byte("ramen")},
+		{key: []byte{0}, value: nil},
+		{
+			key:   []byte("The boxed moved. That was a problem."),
+			value: []byte("The question now was whether or not Peter was going to open it up and look inside to see why it had moved."), // nolint
+		},
+	}
+
+	key := []byte{1, 2}
+	value := []byte{3, 4}
+	const dirty = true
+	const generation = 0
+
+	t.Run("happy path, tries being loaded are same as trie being read", func(t *testing.T) {
+		t.Parallel()
+
+		// hash could be different for keys smaller than 32 and larger than 32 bits.
+		// thus, testing with keys of different sizes.
+		keysToTest := [][]byte{
+			[]byte("This handout will help you understand how paragraphs are formed, how to develop stronger paragraphs."),
+			[]byte("This handout"),
+			[]byte("test"),
+		}
+
+		for _, keyToChild := range keysToTest {
+			trie := NewEmptyTrie()
+
+			for _, test := range keyValue {
+				trie.Put(test.key, test.value)
+			}
+
+			db := newTestDB(t)
+
+			sampleChildTrie := NewTrie(node.NewLeaf(key, value, dirty, generation))
+
+			err := trie.PutChild(keyToChild, sampleChildTrie)
+			require.NoError(t, err)
+
+			err = trie.Store(db)
+			require.NoError(t, err)
+
+			res := NewEmptyTrie()
+
+			err = res.Load(db, trie.MustHash())
+			require.NoError(t, err)
+
+			require.Equal(t, trie.childTries, res.childTries)
+			require.Equal(t, trie.String(), res.String())
+		}
+
+	})
+
+}

lib/trie/helpers_test.go

@@ -20,7 +20,7 @@ type writeCall struct {
 
 var errTest = errors.New("test error")
 
-type Test struct {
+type keyValues struct {
 	key   []byte
 	value []byte
 	op    int