feat(lib/trie): clear fields when node is dirty (#2297)

Author: qdm12

internal/trie/node/dirty.go

@@ -11,6 +11,13 @@ func (b *Branch) IsDirty() bool {
 // SetDirty sets the dirty status to the branch.
 func (b *Branch) SetDirty(dirty bool) {
 	b.Dirty = dirty
+	if dirty {
+		// A node is marked dirty if its key or value is modified.
+		// This means its cached encoding and hash fields are no longer
+		// valid. To improve memory usage, we clear these fields.
+		b.Encoding = nil
+		b.HashDigest = nil
+	}
 }
 
 // IsDirty returns the dirty status of the leaf.
@@ -21,4 +28,11 @@ func (l *Leaf) IsDirty() bool {
 // SetDirty sets the dirty status to the leaf.
 func (l *Leaf) SetDirty(dirty bool) {
 	l.Dirty = dirty
+	if dirty {
+		// A node is marked dirty if its key or value is modified.
+		// This means its cached encoding and hash fields are no longer
+		// valid. To improve memory usage, we clear these fields.
+		l.Encoding = nil
+		l.HashDigest = nil
+	}
 }

internal/trie/node/dirty_test.go

@@ -48,20 +48,40 @@ func Test_Branch_SetDirty(t *testing.T) {
 		expected *Branch
 	}{
 		"not dirty to not dirty": {
-			branch:   &Branch{},
-			expected: &Branch{},
+			branch: &Branch{
+				Encoding:   []byte{1},
+				HashDigest: []byte{1},
+			},
+			expected: &Branch{
+				Encoding:   []byte{1},
+				HashDigest: []byte{1},
+			},
 		},
 		"not dirty to dirty": {
-			branch:   &Branch{},
+			branch: &Branch{
+				Encoding:   []byte{1},
+				HashDigest: []byte{1},
+			},
 			dirty:    true,
 			expected: &Branch{Dirty: true},
 		},
 		"dirty to not dirty": {
-			branch:   &Branch{Dirty: true},
-			expected: &Branch{},
+			branch: &Branch{
+				Encoding:   []byte{1},
+				HashDigest: []byte{1},
+				Dirty:      true,
+			},
+			expected: &Branch{
+				Encoding:   []byte{1},
+				HashDigest: []byte{1},
+			},
 		},
 		"dirty to dirty": {
-			branch:   &Branch{Dirty: true},
+			branch: &Branch{
+				Encoding:   []byte{1},
+				HashDigest: []byte{1},
+				Dirty:      true,
+			},
 			dirty:    true,
 			expected: &Branch{Dirty: true},
 		},
@@ -118,20 +138,40 @@ func Test_Leaf_SetDirty(t *testing.T) {
 		expected *Leaf
 	}{
 		"not dirty to not dirty": {
-			leaf:     &Leaf{},
-			expected: &Leaf{},
+			leaf: &Leaf{
+				Encoding:   []byte{1},
+				HashDigest: []byte{1},
+			},
+			expected: &Leaf{
+				Encoding:   []byte{1},
+				HashDigest: []byte{1},
+			},
 		},
 		"not dirty to dirty": {
-			leaf:     &Leaf{},
+			leaf: &Leaf{
+				Encoding:   []byte{1},
+				HashDigest: []byte{1},
+			},
 			dirty:    true,
 			expected: &Leaf{Dirty: true},
 		},
 		"dirty to not dirty": {
-			leaf:     &Leaf{Dirty: true},
-			expected: &Leaf{},
+			leaf: &Leaf{
+				Encoding:   []byte{1},
+				HashDigest: []byte{1},
+				Dirty:      true,
+			},
+			expected: &Leaf{
+				Encoding:   []byte{1},
+				HashDigest: []byte{1},
+			},
 		},
 		"dirty to dirty": {
-			leaf:     &Leaf{Dirty: true},
+			leaf: &Leaf{
+				Encoding:   []byte{1},
+				HashDigest: []byte{1},
+				Dirty:      true,
+			},
 			dirty:    true,
 			expected: &Leaf{Dirty: true},
 		},

internal/trie/node/leaf.go

@@ -68,5 +68,4 @@ func bytesToString(b []byte) (s string) {
 	default:
 		return fmt.Sprintf("0x%x...%x", b[:8], b[len(b)-8:])
 	}
-
 }

lib/trie/trie.go

@@ -366,7 +366,7 @@ func (t *Trie) insertInLeaf(parentLeaf *node.Leaf, key,
 			// Move the current leaf parent as a child to the new branch.
 			childIndex := parentLeafKey[commonPrefixLength]
 			parentLeaf.Key = parentLeaf.Key[commonPrefixLength+1:]
-			parentLeaf.Dirty = true
+			parentLeaf.SetDirty(true)
 			newBranchParent.Children[childIndex] = parentLeaf
 		}
 
@@ -440,7 +440,6 @@ func (t *Trie) insertInBranch(parentBranch *node.Branch, key, value []byte) (new
 	oldParentIndex := parentBranch.Key[commonPrefixLength]
 	remainingOldParentKey := parentBranch.Key[commonPrefixLength+1:]
 
-	parentBranch.Dirty = true
 	parentBranch.Key = remainingOldParentKey
 	parentBranch.Generation = t.generation
 	newParentBranch.Children[oldParentIndex] = parentBranch