fix(trie): PopulateMerkleValues behavior

- Include Merkle value of the parent node passed as argument
- Calculate missing Merkle values if needed
- Add unit test

Author: qdm12

lib/trie/database.go

@@ -187,21 +187,44 @@ func (t *Trie) loadNode(db Database, n *Node) error {
 
 // PopulateMerkleValues writes the Merkle value of each children of the node given
 // as keys to the map merkleValues.
-func (t *Trie) PopulateMerkleValues(n *Node, merkleValues map[string]struct{}) {
-	if n.Kind() != node.Branch {
-		return
+func (t *Trie) PopulateMerkleValues(n *Node,
+	merkleValues map[string]struct{}) (err error) {
+	if n == nil {
+		return nil
 	}
 
-	branch := n
-	for _, child := range branch.Children {
-		if child == nil {
-			continue
+	merkleValue := n.MerkleValue
+	if len(merkleValue) == 0 {
+		// Compute and cache node Merkle value if it is absent.
+		if n == t.root {
+			merkleValue, err = n.CalculateRootMerkleValue()
+			if err != nil {
+				return fmt.Errorf("calculating Merkle value for root node: %w", err)
+			}
+		} else {
+			merkleValue, err = n.CalculateMerkleValue()
+			if err != nil {
+				return fmt.Errorf("calculating Merkle value for node: %w", err)
+			}
 		}
+	}
 
-		merkleValues[string(child.MerkleValue)] = struct{}{}
+	merkleValues[string(merkleValue)] = struct{}{}
+
+	if n.Kind() == node.Leaf {
+		return nil
+	}
 
-		t.PopulateMerkleValues(child, merkleValues)
+	branch := n
+	for _, child := range branch.Children {
+		err = t.PopulateMerkleValues(child, merkleValues)
+		if err != nil {
+			// Note: do not wrap error since this is recursive.
+			return err
+		}
 	}
+
+	return nil
 }
 
 // GetFromDB retrieves a value at the given key from the trie using the database.

lib/trie/database_test.go

@@ -158,7 +158,100 @@ func Test_Trie_WriteDirty_ClearPrefix(t *testing.T) {
 	assert.Equal(t, trie.String(), trieFromDB.String())
 }
 
-func Test_Trie_GetFromDB(t *testing.T) {
+func Test_PopulateMerkleValues(t *testing.T) {
+	t.Parallel()
+
+	someNode := &Node{Key: []byte{1}, SubValue: []byte{2}}
+
+	testCases := map[string]struct {
+		trie         *Trie
+		node         *Node
+		merkleValues map[string]struct{}
+		errSentinel  error
+		errMessage   string
+	}{
+		"nil node": {
+			trie:         &Trie{},
+			merkleValues: map[string]struct{}{},
+		},
+		"leaf node": {
+			trie: &Trie{},
+			node: &Node{MerkleValue: []byte("a")},
+			merkleValues: map[string]struct{}{
+				"a": {},
+			},
+		},
+		"leaf node without Merkle value": {
+			trie: &Trie{},
+			node: &Node{Key: []byte{1}, SubValue: []byte{2}},
+			merkleValues: map[string]struct{}{
+				"A\x01\x04\x02": {},
+			},
+		},
+		"root leaf node without Merkle value": {
+			trie: &Trie{
+				root: someNode,
+			},
+			node: someNode,
+			merkleValues: map[string]struct{}{
+				"`Qm\v\xb6\xe1\xbb\xfb\x12\x93\xf1\xb2v\xea\x95\x05\xe9\xf4\xa4\xe7ُb\r\x05\x11^\v\x85'J\xe1": {}, //nolint:lll
+			},
+		},
+		"branch node": {
+			trie: &Trie{},
+			node: &Node{
+				MerkleValue: []byte("a"),
+				Children: padRightChildren([]*Node{
+					{MerkleValue: []byte("b")},
+				}),
+			},
+			merkleValues: map[string]struct{}{
+				"a": {},
+				"b": {},
+			},
+		},
+		"nested branch node": {
+			trie: &Trie{},
+			node: &Node{
+				MerkleValue: []byte("a"),
+				Children: padRightChildren([]*Node{
+					{MerkleValue: []byte("b")},
+					{
+						MerkleValue: []byte("c"),
+						Children: padRightChildren([]*Node{
+							{MerkleValue: []byte("d")},
+						}),
+					},
+				}),
+			},
+			merkleValues: map[string]struct{}{
+				"a": {},
+				"b": {},
+				"c": {},
+				"d": {},
+			},
+		},
+	}
+
+	for name, testCase := range testCases {
+		testCase := testCase
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			merkleValues := make(map[string]struct{})
+
+			err := testCase.trie.PopulateMerkleValues(testCase.node, merkleValues)
+
+			assert.ErrorIs(t, err, testCase.errSentinel)
+			if testCase.errSentinel != nil {
+				assert.EqualError(t, err, testCase.errMessage)
+			}
+			assert.Equal(t, testCase.merkleValues, merkleValues)
+		})
+	}
+}
+
+func Test_GetFromDB(t *testing.T) {
 	t.Parallel()
 
 	const size = 1000