Change ensureMerkleValueIsCalculated to check for root node

ChainSafe · Nov 23, 2022 · de33406 · de33406
1 parent e432967
commit de33406
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 31 deletions.
diff --git a/lib/trie/trie.go b/lib/trie/trie.go
@@ -104,8 +104,7 @@ func (t *Trie) prepForMutation(currentNode *Node,
 
 func (t *Trie) registerDeletedMerkleValue(node *Node,
 	pendingDeletedMerkleValues map[string]struct{}) (err error) {
-	isRoot := node == t.root
-	err = ensureMerkleValueIsCalculated(node, isRoot)
+	err = t.ensureMerkleValueIsCalculated(node)
 	if err != nil {
 		return fmt.Errorf("ensuring Merkle value is calculated: %w", err)
 	}
@@ -1025,8 +1024,7 @@ func (t *Trie) ClearPrefix(prefixLE []byte) (err error) {
 	}()
 
 	if len(prefixLE) == 0 {
-		const isRoot = true
-		err = ensureMerkleValueIsCalculated(t.root, isRoot)
+		err = t.ensureMerkleValueIsCalculated(t.root)
 		if err != nil {
 			return fmt.Errorf("ensuring Merkle values are calculated: %w", err)
 		}
@@ -1057,8 +1055,7 @@ func (t *Trie) clearPrefixAtNode(parent *Node, prefix []byte,
 	}
 
 	if bytes.HasPrefix(parent.PartialKey, prefix) {
-		isRoot := parent == t.root
-		err = ensureMerkleValueIsCalculated(parent, isRoot)
+		err = t.ensureMerkleValueIsCalculated(parent)
 		if err != nil {
 			nodesRemoved = 0
 			return parent, nodesRemoved, fmt.Errorf("ensuring Merkle values are calculated: %w", err)
@@ -1375,12 +1372,12 @@ func (t *Trie) handleDeletion(branch *Node, key []byte,
 // to ensure the parent node and all its descendant nodes have their Merkle
 // value computed and ready to be used. This has a close to zero performance
 // impact if the parent node Merkle value is already computed.
-func ensureMerkleValueIsCalculated(parent *Node, isRoot bool) (err error) {
+func (t *Trie) ensureMerkleValueIsCalculated(parent *Node) (err error) {
 	if parent == nil {
 		return nil
 	}
 
-	if isRoot {
+	if parent == t.root {
 		_, err = parent.CalculateRootMerkleValue()
 		if err != nil {
 			return fmt.Errorf("calculating Merkle value of root node: %w", err)

diff --git a/lib/trie/trie_test.go b/lib/trie/trie_test.go
@@ -317,8 +317,8 @@ func Test_Trie_registerDeletedMerkleValue(t *testing.T) {
 			},
 			expectedTrie: Trie{
 				root: &Node{
-					PartialKey: []byte{1},
-					SubValue:   []byte{2},
+					PartialKey:   []byte{1},
+					StorageValue: []byte{2},
 					MerkleValue: []byte{
 						0x60, 0x51, 0x6d, 0x0b, 0xb6, 0xe1, 0xbb, 0xfb,
 						0x12, 0x93, 0xf1, 0xb2, 0x76, 0xea, 0x95, 0x05,
@@ -4066,23 +4066,44 @@ func Test_Trie_handleDeletion(t *testing.T) {
 	}
 }
 
-func Test_ensureMerkleValueIsCalculated(t *testing.T) {
+func Test_Trie_ensureMerkleValueIsCalculated(t *testing.T) {
 	t.Parallel()
 
+	node := &Node{
+		PartialKey:   []byte{1},
+		StorageValue: []byte{2},
+	}
+
+	nodeWithEncodingMerkleValue := &Node{
+		PartialKey:   []byte{1},
+		StorageValue: []byte{2},
+		MerkleValue:  []byte{3},
+	}
+
+	nodeWithHashMerkleValue := &Node{
+		PartialKey:   []byte{1},
+		StorageValue: []byte{2},
+		MerkleValue: []byte{
+			1, 2, 3, 4, 5, 6, 7, 8,
+			1, 2, 3, 4, 5, 6, 7, 8,
+			1, 2, 3, 4, 5, 6, 7, 8,
+			1, 2, 3, 4, 5, 6, 7, 8},
+	}
+
 	testCases := map[string]struct {
+		trie         Trie
 		parent       *Node
-		isRoot       bool
 		errSentinel  error
 		errMessage   string
 		expectedNode *Node
+		expectedTrie Trie
 	}{
 		"nil parent": {},
 		"root node without Merkle value": {
-			parent: &Node{
-				PartialKey:   []byte{1},
-				StorageValue: []byte{2},
+			trie: Trie{
+				root: node,
 			},
-			isRoot: true,
+			parent: node,
 			expectedNode: &Node{
 				PartialKey:   []byte{1},
 				StorageValue: []byte{2},
@@ -4092,14 +4113,15 @@ func Test_ensureMerkleValueIsCalculated(t *testing.T) {
 					0xe9, 0xf4, 0xa4, 0xe7, 0xd9, 0x8f, 0x62, 0xd,
 					0x5, 0x11, 0x5e, 0xb, 0x85, 0x27, 0x4a, 0xe1},
 			},
+			expectedTrie: Trie{
+				root: node,
+			},
 		},
 		"root node with inlined Merkle value": {
-			parent: &Node{
-				PartialKey:   []byte{1},
-				StorageValue: []byte{2},
-				MerkleValue:  []byte{3},
+			trie: Trie{
+				root: nodeWithEncodingMerkleValue,
 			},
-			isRoot: true,
+			parent: nodeWithEncodingMerkleValue,
 			expectedNode: &Node{
 				PartialKey:   []byte{1},
 				StorageValue: []byte{2},
@@ -4109,18 +4131,15 @@ func Test_ensureMerkleValueIsCalculated(t *testing.T) {
 					0xe9, 0xf4, 0xa4, 0xe7, 0xd9, 0x8f, 0x62, 0xd,
 					0x5, 0x11, 0x5e, 0xb, 0x85, 0x27, 0x4a, 0xe1},
 			},
+			expectedTrie: Trie{
+				root: nodeWithEncodingMerkleValue,
+			},
 		},
 		"root node with hash Merkle value": {
-			parent: &Node{
-				PartialKey:   []byte{1},
-				StorageValue: []byte{2},
-				MerkleValue: []byte{
-					1, 2, 3, 4, 5, 6, 7, 8,
-					1, 2, 3, 4, 5, 6, 7, 8,
-					1, 2, 3, 4, 5, 6, 7, 8,
-					1, 2, 3, 4, 5, 6, 7, 8},
+			trie: Trie{
+				root: nodeWithHashMerkleValue,
 			},
-			isRoot: true,
+			parent: nodeWithHashMerkleValue,
 			expectedNode: &Node{
 				PartialKey:   []byte{1},
 				StorageValue: []byte{2},
@@ -4130,6 +4149,9 @@ func Test_ensureMerkleValueIsCalculated(t *testing.T) {
 					1, 2, 3, 4, 5, 6, 7, 8,
 					1, 2, 3, 4, 5, 6, 7, 8},
 			},
+			expectedTrie: Trie{
+				root: nodeWithHashMerkleValue,
+			},
 		},
 		"non root node without Merkle value": {
 			parent: &Node{
@@ -4161,14 +4183,15 @@ func Test_ensureMerkleValueIsCalculated(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			t.Parallel()
 
-			err := ensureMerkleValueIsCalculated(testCase.parent, testCase.isRoot)
+			err := testCase.trie.ensureMerkleValueIsCalculated(testCase.parent)
 
 			checkMerkleValuesAreSet(t, testCase.parent)
 			assert.ErrorIs(t, err, testCase.errSentinel)
 			if testCase.errSentinel != nil {
 				assert.EqualError(t, err, testCase.errMessage)
 			}
 			assert.Equal(t, testCase.expectedNode, testCase.parent)
+			assert.Equal(t, testCase.expectedTrie, testCase.trie)
 		})
 	}
 }