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

dot/state/base_test.go

@@ -35,7 +35,7 @@ func TestTrie_StoreAndLoadFromDB(t *testing.T) {
 	expected := tt.MustHash()
 
 	tt = trie.NewEmptyTrie()
-	err = tt.Load(db, encroot)
+	err = tt.Load(db, encroot.ToBytes())
 	require.NoError(t, err)
 	require.Equal(t, expected, tt.MustHash())
 }

dot/state/storage.go

@@ -142,7 +142,7 @@ func (s *StorageState) TrieState(root *common.Hash) (*rtstorage.TrieState, error
 // LoadFromDB loads an encoded trie from the DB where the key is `root`
 func (s *StorageState) LoadFromDB(root common.Hash) (*trie.Trie, error) {
 	t := trie.NewEmptyTrie()
-	err := t.Load(s.db, root)
+	err := t.Load(s.db, root.ToBytes())
 	if err != nil {
 		return nil, err
 	}

dot/state/storage_test.go

@@ -11,6 +11,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"
@@ -227,7 +228,12 @@ func TestGetStorageChildAndGetStorageFromChild(t *testing.T) {
 		"0",
 	))
 
-	testChildTrie := trie.NewEmptyTrie()
+	key := []byte{1, 2}
+	value := []byte{3, 4}
+	const dirty = true
+	const generation = 0
+	testChildTrie := trie.NewTrie(node.NewLeaf(key, value, dirty, generation))
+
 	testChildTrie.Put([]byte("keyInsidechild"), []byte("voila"))
 
 	err = genTrie.PutChild([]byte("keyToChild"), testChildTrie)
@@ -236,6 +242,7 @@ func TestGetStorageChildAndGetStorageFromChild(t *testing.T) {
 	triesGauge := NewMockGauge(ctrl)
 	triesGauge.EXPECT().Inc().Times(2)
 	triesGauge.EXPECT().Set(0.00)
+
 	tries := &Tries{
 		rootToTrie: make(map[common.Hash]*trie.Trie),
 		triesGauge: triesGauge,
@@ -269,7 +276,7 @@ func TestGetStorageChildAndGetStorageFromChild(t *testing.T) {
 	_, err = storage.GetStorageChild(&rootHash, []byte("keyToChild"))
 	require.NoError(t, err)
 
-	value, err := storage.GetStorageFromChild(&rootHash, []byte("keyToChild"), []byte("keyInsidechild"))
+	value, err = storage.GetStorageFromChild(&rootHash, []byte("keyToChild"), []byte("keyInsidechild"))
 	require.NoError(t, err)
 
 	require.Equal(t, []byte("voila"), value)

lib/trie/child_storage.go

@@ -16,17 +16,20 @@ 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 main trie. K is the child storage key
+// associated to the child trie, and V is the Merkle value of its corresponding child trie.
+// To get merkle value of a node, we get value of the node. If the value is shorter than
+// 32 bits, we use it as merkle value. Otherwise, we use it's blake2b hash as merkle value.
 func (t *Trie) PutChild(keyToChild []byte, child *Trie) error {
-	childHash, err := child.Hash()
+	_, hash, err := child.root.EncodeAndHash()
 	if err != nil {
 		return err
 	}
 
 	key := append(ChildStorageKeyPrefix, keyToChild...)
-	value := [32]byte(childHash)
 
-	t.Put(key, value[:])
-	t.childTries[childHash] = child
+	t.Put(key, hash)
+	t.childTries[common.BytesToHash(hash)] = child
 	return nil
 }
 

lib/trie/database.go

@@ -145,24 +145,23 @@ func (t *Trie) loadProof(proofHashToNode map[string]Node, n Node) {
 
 // Load reconstructs the trie from the database from the given root hash.
 // It is used when restarting the node to load the current state trie.
-func (t *Trie) Load(db chaindb.Database, rootHash common.Hash) error {
-	if rootHash == EmptyHash {
+func (t *Trie) Load(db chaindb.Database, rootHashBytes []byte) error {
+	if len(rootHashBytes) == 0 {
 		t.root = nil
 		return nil
 	}
 
-	rootHashBytes := rootHash[:]
-
 	encodedNode, err := db.Get(rootHashBytes)
 	if err != nil {
-		return fmt.Errorf("failed to find root key %s: %w", rootHash, err)
+		return fmt.Errorf("failed to find root key 0x%x: %w", rootHashBytes, err)
 	}
 
 	reader := bytes.NewReader(encodedNode)
 	root, err := node.Decode(reader)
 	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 +184,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,11 +209,13 @@ 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))
+		// TODO: Tests this error
+		err := childTrie.Load(db, value)
 		if err != nil {
 			return fmt.Errorf("failed to load child trie with root hash=0x%x: %w", value, err)
 		}
 
+		// TODO: Test this error
 		err = t.PutChild(value, childTrie)
 		if err != nil {
 			return fmt.Errorf("failed to insert child trie with root hash=0x%x into main trie: %w",
@@ -377,8 +379,8 @@ func (t *Trie) writeDirty(db chaindb.Batch, n Node) error {
 			n.GetHash(), err)
 	}
 
-	if n == t.root {
-		// hash root node even if its encoding is under 32 bytes
+	// hash root node even if its encoding is under 32 bytes
+	if n == t.root && len(encoding) < 32 {
 		encodingDigest, err := common.Blake2bHash(encoding)
 		if err != nil {
 			return fmt.Errorf("cannot hash root node encoding: %w", err)

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"
 )
@@ -63,9 +64,10 @@ func TestTrie_DatabaseStoreAndLoad(t *testing.T) {
 		require.NoError(t, err)
 
 		res := NewEmptyTrie()
-		err = res.Load(db, trie.MustHash())
+		err = res.Load(db, trie.MustHash().ToBytes())
 		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)
@@ -134,7 +136,7 @@ func TestTrie_WriteDirty_Put(t *testing.T) {
 		require.NoError(t, err)
 
 		res := NewEmptyTrie()
-		err = res.Load(db, trie.MustHash())
+		err = res.Load(db, trie.MustHash().ToBytes())
 		require.NoError(t, err)
 		require.Equal(t, trie.MustHash(), res.MustHash())
 
@@ -204,7 +206,7 @@ func TestTrie_WriteDirty_PutReplace(t *testing.T) {
 		}
 
 		res := NewEmptyTrie()
-		err := res.Load(db, trie.MustHash())
+		err := res.Load(db, trie.MustHash().ToBytes())
 		require.NoError(t, err)
 		require.Equal(t, trie.MustHash(), res.MustHash())
 
@@ -264,7 +266,7 @@ func TestTrie_WriteDirty_Delete(t *testing.T) {
 			require.NoError(t, err)
 
 			res := NewEmptyTrie()
-			err = res.Load(db, trie.MustHash())
+			err = res.Load(db, trie.MustHash().ToBytes())
 			require.NoError(t, err)
 			require.Equal(t, trie.MustHash(), res.MustHash())
 
@@ -329,7 +331,7 @@ func TestTrie_WriteDirty_ClearPrefix(t *testing.T) {
 		require.NoError(t, err)
 
 		res := NewEmptyTrie()
-		err = res.Load(db, trie.MustHash())
+		err = res.Load(db, trie.MustHash().ToBytes())
 		require.NoError(t, err)
 
 		require.Equal(t, trie.MustHash(), res.MustHash())
@@ -387,3 +389,58 @@ func TestTrie_GetFromDB(t *testing.T) {
 		}
 	}
 }
+
+func TestStoreAndLoadWithChildTries(t *testing.T) {
+	testCase := []Test{
+		{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) {
+		// 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 testCase {
+				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.root.GetHash())
+			require.NoError(t, err)
+
+			require.Equal(t, trie.childTries, res.childTries)
+		}
+
+	})
+
+}

lib/trie/proof.go

@@ -37,7 +37,7 @@ func GenerateProof(root []byte, keys [][]byte, db chaindb.Database) ([][]byte, e
 	trackedProofs := make(map[string][]byte)
 
 	proofTrie := NewEmptyTrie()
-	if err := proofTrie.Load(db, common.BytesToHash(root)); err != nil {
+	if err := proofTrie.Load(db, root); err != nil {
 		return nil, err
 	}
 

lib/trie/trie_endtoend_test.go

@@ -22,7 +22,6 @@ import (
 
 	"github.com/ChainSafe/gossamer/internal/trie/codec"
 	"github.com/ChainSafe/gossamer/internal/trie/node"
-	"github.com/ChainSafe/gossamer/lib/common"
 )
 
 const (
@@ -397,7 +396,7 @@ func TestTrieDiff(t *testing.T) {
 	}
 
 	dbTrie := NewEmptyTrie()
-	err = dbTrie.Load(storageDB, common.BytesToHash(newTrie.root.GetHash()))
+	err = dbTrie.Load(storageDB, newTrie.root.GetHash())
 	require.NoError(t, err)
 }