From d503984d90e6da5c69a65ebd78ce7535d6f25332 Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Tue, 4 Oct 2022 13:27:50 +0000
Subject: [PATCH] No Merkle value computation in function - Update comment that
 all Merkle values are expected to be set - Panic if merkle value is not set -
 Avoid leaving the map in a bad state

---
 dot/state/offline_pruner.go |  5 +----
 lib/trie/database.go        | 39 +++++++++++--------------------------
 lib/trie/database_test.go   | 39 +++++++++++--------------------------
 3 files changed, 23 insertions(+), 60 deletions(-)

diff --git a/dot/state/offline_pruner.go b/dot/state/offline_pruner.go
index 5c591e5483..72d400260b 100644
--- a/dot/state/offline_pruner.go
+++ b/dot/state/offline_pruner.go
@@ -121,10 +121,7 @@ func (p *OfflinePruner) SetBloomFilter() (err error) {
 			return err
 		}
 
-		err = tr.PopulateMerkleValues(tr.RootNode(), merkleValues)
-		if err != nil {
-			return fmt.Errorf("populating Merkle values from trie: %w", err)
-		}
+		trie.PopulateMerkleValues(tr.RootNode(), merkleValues)
 
 		// get parent header of current block
 		header, err = p.blockState.GetHeader(header.ParentHash)
diff --git a/lib/trie/database.go b/lib/trie/database.go
index bd0501d34c..e5c59f8897 100644
--- a/lib/trie/database.go
+++ b/lib/trie/database.go
@@ -185,46 +185,29 @@ func (t *Trie) loadNode(db Database, n *Node) error {
 	return nil
 }
 
-// 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{}) (err error) {
+// PopulateMerkleValues writes the Merkle values of the node given and of
+// all its descendant nodes as keys to the map merkleValues.
+// It is assumed the node and its descendant nodes have their Merkle value already
+// computed.
+func PopulateMerkleValues(n *Node, merkleValues map[string]struct{}) {
 	if n == nil {
-		return nil
+		return
 	}
 
-	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)
-			}
-		}
+	if len(n.MerkleValue) == 0 {
+		panic(fmt.Sprintf("node with key 0x%x has no Merkle value computed", n.Key))
 	}
 
-	merkleValues[string(merkleValue)] = struct{}{}
+	merkleValues[string(n.MerkleValue)] = struct{}{}
 
 	if n.Kind() == node.Leaf {
-		return nil
+		return
 	}
 
 	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
-		}
+		PopulateMerkleValues(child, merkleValues)
 	}
-
-	return nil
 }
 
 // GetFromDB retrieves a value at the given key from the trie using the database.
diff --git a/lib/trie/database_test.go b/lib/trie/database_test.go
index cfca9497dd..63a2883d1a 100644
--- a/lib/trie/database_test.go
+++ b/lib/trie/database_test.go
@@ -161,44 +161,25 @@ func Test_Trie_WriteDirty_ClearPrefix(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
+		panicValue   interface{}
 	}{
 		"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": {},
-			},
+			node:       &Node{Key: []byte{1}, SubValue: []byte{2}},
+			panicValue: "node with key 0x01 has no Merkle value computed",
 		},
 		"branch node": {
-			trie: &Trie{},
 			node: &Node{
 				MerkleValue: []byte("a"),
 				Children: padRightChildren([]*Node{
@@ -211,7 +192,6 @@ func Test_PopulateMerkleValues(t *testing.T) {
 			},
 		},
 		"nested branch node": {
-			trie: &Trie{},
 			node: &Node{
 				MerkleValue: []byte("a"),
 				Children: padRightChildren([]*Node{
@@ -240,12 +220,15 @@ func Test_PopulateMerkleValues(t *testing.T) {
 
 			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)
+			if testCase.panicValue != nil {
+				assert.PanicsWithValue(t, testCase.panicValue, func() {
+					PopulateMerkleValues(testCase.node, merkleValues)
+				})
+				return
 			}
+
+			PopulateMerkleValues(testCase.node, merkleValues)
+
 			assert.Equal(t, testCase.merkleValues, merkleValues)
 		})
 	}