Merge updateBranch and insertInBranch

ChainSafe · Feb 16, 2022 · 52cb3f3 · 52cb3f3
1 parent 9757382
commit 52cb3f3
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 17 deletions.
diff --git a/lib/trie/trie.go b/lib/trie/trie.go
@@ -334,19 +334,6 @@ func (t *Trie) insert(parent Node, key []byte, value Node) (newParent Node) {
 	}
 }
 
-func (t *Trie) insertInBranch(parentBranch *node.Branch, key []byte,
-	value Node) (newParent Node) {
-	newParent = t.updateBranch(parentBranch, key, value)
-
-	if newParent.IsDirty() {
-		// the older parent branch might had been pushed down the trie
-		// under the new parent branch, so mark it dirty.
-		parentBranch.SetDirty(true)
-	}
-
-	return newParent
-}
-
 func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key []byte,
 	value Node) (newParent Node) {
 	newValue := value.(*node.Leaf).Value
@@ -402,7 +389,7 @@ func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key []byte,
 	return newBranchParent
 }
 
-func (t *Trie) updateBranch(parentBranch *node.Branch, key []byte, value Node) (newParent Node) {
+func (t *Trie) insertInBranch(parentBranch *node.Branch, key []byte, value Node) (newParent Node) {
 	if bytes.Equal(key, parentBranch.Key) {
 		parentBranch.SetDirty(true)
 		parentBranch.Value = value.GetValue()
@@ -441,6 +428,7 @@ func (t *Trie) updateBranch(parentBranch *node.Branch, key []byte, value Node) (
 		Generation: t.generation,
 		Dirty:      true,
 	}
+	parentBranch.SetDirty(true)
 
 	oldParentIndex := parentBranch.Key[commonPrefixLength]
 	remainingOldParentKey := parentBranch.Key[commonPrefixLength+1:]
@@ -454,7 +442,6 @@ func (t *Trie) updateBranch(parentBranch *node.Branch, key []byte, value Node) (
 		newParentBranch.Children[childIndex] = t.insert(nil, remainingKey, value)
 	}
 
-	newParentBranch.SetDirty(true)
 	return newParentBranch
 }
 

diff --git a/lib/trie/trie_test.go b/lib/trie/trie_test.go
@@ -1142,7 +1142,7 @@ func Test_Trie_insert(t *testing.T) {
 	}
 }
 
-func Test_Trie_updateBranch(t *testing.T) {
+func Test_Trie_insertInBranch(t *testing.T) {
 	t.Parallel()
 
 	testCases := map[string]struct {
@@ -1280,6 +1280,7 @@ func Test_Trie_updateBranch(t *testing.T) {
 					&node.Branch{
 						Key:   []byte{},
 						Value: []byte{5},
+						Dirty: true,
 						Children: [16]node.Node{
 							&node.Leaf{Key: []byte{1}},
 						},
@@ -1311,6 +1312,7 @@ func Test_Trie_updateBranch(t *testing.T) {
 					&node.Branch{
 						Key:   []byte{3},
 						Value: []byte{5},
+						Dirty: true,
 						Children: [16]node.Node{
 							&node.Leaf{Key: []byte{1}},
 						},
@@ -1343,6 +1345,7 @@ func Test_Trie_updateBranch(t *testing.T) {
 					&node.Branch{
 						Key:   []byte{},
 						Value: []byte{5},
+						Dirty: true,
 						Children: [16]node.Node{
 							&node.Leaf{Key: []byte{1}},
 						},
@@ -1359,7 +1362,7 @@ func Test_Trie_updateBranch(t *testing.T) {
 
 			trie := new(Trie)
 
-			newNode := trie.updateBranch(testCase.parent, testCase.key, testCase.value)
+			newNode := trie.insertInBranch(testCase.parent, testCase.key, testCase.value)
 
 			assert.Equal(t, testCase.newNode, newNode)
 			assert.Equal(t, new(Trie), trie) // check no mutation