From 7e7d1a7bcf78d5a69d3196f3989de5b24536da62 Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Wed, 16 Feb 2022 18:40:22 +0000
Subject: [PATCH] feat(lib/trie): clear fields when node is dirty

---
 internal/trie/node/dirty.go      | 14 +++++++
 internal/trie/node/dirty_test.go | 64 ++++++++++++++++++++++++++------
 internal/trie/node/leaf.go       |  1 -
 lib/trie/trie.go                 |  3 +-
 4 files changed, 67 insertions(+), 15 deletions(-)

diff --git a/internal/trie/node/dirty.go b/internal/trie/node/dirty.go
index 7922139b18..c4d13f3b61 100644
--- a/internal/trie/node/dirty.go
+++ b/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
+	}
 }
diff --git a/internal/trie/node/dirty_test.go b/internal/trie/node/dirty_test.go
index c2849b0932..62aedb2616 100644
--- a/internal/trie/node/dirty_test.go
+++ b/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},
 		},
diff --git a/internal/trie/node/leaf.go b/internal/trie/node/leaf.go
index b5a5e02b2a..47a6dd0325 100644
--- a/internal/trie/node/leaf.go
+++ b/internal/trie/node/leaf.go
@@ -71,5 +71,4 @@ func bytesToString(b []byte) (s string) {
 	default:
 		return fmt.Sprintf("0x%x...%x", b[:8], b[len(b)-8:])
 	}
-
 }
diff --git a/lib/trie/trie.go b/lib/trie/trie.go
index e6df9702a8..d3a183ddd8 100644
--- a/lib/trie/trie.go
+++ b/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