feat(trie): finer deep copy of nodes (#2384)

- Do not copy cached fields when not needed - Do not copy key field when not needed - Do not copy value field when not needed
ChainSafe · Mar 28, 2022 · bd6d8e4 · bd6d8e4
1 parent 50652e9
commit bd6d8e4
Show file tree

Hide file tree

Showing 5 changed files with 182 additions and 60 deletions.
diff --git a/internal/trie/node/copy.go b/internal/trie/node/copy.go
@@ -3,76 +3,134 @@
 
 package node
 
+var (
+	// DefaultCopySettings contains the following copy settings:
+	// - children are NOT deep copied recursively
+	// - the HashDigest field is left empty on the copy
+	// - the Encoding field is left empty on the copy
+	// - the key field is deep copied
+	// - the value field is deep copied
+	DefaultCopySettings = CopySettings{
+		CopyKey:   true,
+		CopyValue: true,
+	}
+
+	// DeepCopySettings returns the following copy settings:
+	// - children are deep copied recursively
+	// - the HashDigest field is deep copied
+	// - the Encoding field is deep copied
+	// - the key field is deep copied
+	// - the value field is deep copied
+	DeepCopySettings = CopySettings{
+		CopyChildren: true,
+		CopyCached:   true,
+		CopyKey:      true,
+		CopyValue:    true,
+	}
+)
+
+// CopySettings contains settings to configure the deep copy
+// of a node.
+type CopySettings struct {
+	// CopyChildren can be set to true to recursively deep copy the eventual
+	// children of the node. This is false by default and should only be used
+	// in tests since it is quite inefficient.
+	CopyChildren bool
+	// CopyCached can be set to true to deep copy the cached digest
+	// and encoding fields on the current node copied.
+	// This is false by default because in production, a node is copied
+	// when it is about to be mutated, hence making its cached fields
+	// no longer valid.
+	CopyCached bool
+	// CopyKey can be set to true to deep copy the key field of
+	// the node. This is useful when false if the key is about to
+	// be assigned after the Copy operation, to save a memory operation.
+	CopyKey bool
+	// CopyValue can be set to true to deep copy the value field of
+	// the node. This is useful when false if the value is about to
+	// be assigned after the Copy operation, to save a memory operation.
+	CopyValue bool
+}
+
 // Copy deep copies the branch.
 // Setting copyChildren to true will deep copy
 // children as well.
-func (b *Branch) Copy(copyChildren bool) Node {
+func (b *Branch) Copy(settings CopySettings) Node {
 	cpy := &Branch{
 		Dirty:      b.Dirty,
 		Generation: b.Generation,
 	}
 
-	if copyChildren {
+	if settings.CopyChildren {
+		// Copy all fields of children if we deep copy children
+		childSettings := settings
+		childSettings.CopyKey = true
+		childSettings.CopyValue = true
+		childSettings.CopyCached = true
 		for i, child := range b.Children {
 			if child == nil {
 				continue
 			}
-			cpy.Children[i] = child.Copy(copyChildren)
+			cpy.Children[i] = child.Copy(childSettings)
 		}
 	} else {
 		cpy.Children = b.Children // copy interface pointers only
 	}
 
-	if b.Key != nil {
+	if settings.CopyKey && b.Key != nil {
 		cpy.Key = make([]byte, len(b.Key))
 		copy(cpy.Key, b.Key)
 	}
 
 	// nil and []byte{} are encoded differently, watch out!
-	if b.Value != nil {
+	if settings.CopyValue && b.Value != nil {
 		cpy.Value = make([]byte, len(b.Value))
 		copy(cpy.Value, b.Value)
 	}
 
-	if b.HashDigest != nil {
-		cpy.HashDigest = make([]byte, len(b.HashDigest))
-		copy(cpy.HashDigest, b.HashDigest)
-	}
+	if settings.CopyCached {
+		if b.HashDigest != nil {
+			cpy.HashDigest = make([]byte, len(b.HashDigest))
+			copy(cpy.HashDigest, b.HashDigest)
+		}
 
-	if b.Encoding != nil {
-		cpy.Encoding = make([]byte, len(b.Encoding))
-		copy(cpy.Encoding, b.Encoding)
+		if b.Encoding != nil {
+			cpy.Encoding = make([]byte, len(b.Encoding))
+			copy(cpy.Encoding, b.Encoding)
+		}
 	}
 
 	return cpy
 }
 
 // Copy deep copies the leaf.
-func (l *Leaf) Copy(_ bool) Node {
+func (l *Leaf) Copy(settings CopySettings) Node {
 	cpy := &Leaf{
 		Dirty:      l.Dirty,
 		Generation: l.Generation,
 	}
 
-	if l.Key != nil {
+	if settings.CopyKey && l.Key != nil {
 		cpy.Key = make([]byte, len(l.Key))
 		copy(cpy.Key, l.Key)
 	}
 
 	// nil and []byte{} are encoded differently, watch out!
-	if l.Value != nil {
+	if settings.CopyValue && l.Value != nil {
 		cpy.Value = make([]byte, len(l.Value))
 		copy(cpy.Value, l.Value)
 	}
 
-	if l.HashDigest != nil {
-		cpy.HashDigest = make([]byte, len(l.HashDigest))
-		copy(cpy.HashDigest, l.HashDigest)
-	}
+	if settings.CopyCached {
+		if l.HashDigest != nil {
+			cpy.HashDigest = make([]byte, len(l.HashDigest))
+			copy(cpy.HashDigest, l.HashDigest)
+		}
 
-	if l.Encoding != nil {
-		cpy.Encoding = make([]byte, len(l.Encoding))
-		copy(cpy.Encoding, l.Encoding)
+		if l.Encoding != nil {
+			cpy.Encoding = make([]byte, len(l.Encoding))
+			copy(cpy.Encoding, l.Encoding)
+		}
 	}
 
 	return cpy

diff --git a/internal/trie/node/copy_test.go b/internal/trie/node/copy_test.go
@@ -4,6 +4,7 @@
 package node
 
 import (
+	"reflect"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -12,8 +13,7 @@ import (
 
 func testForSliceModif(t *testing.T, original, copied []byte) {
 	t.Helper()
-	require.Equal(t, len(original), len(copied))
-	if len(copied) == 0 {
+	if !reflect.DeepEqual(original, copied) || len(copied) == 0 {
 		// cannot test for modification
 		return
 	}
@@ -26,7 +26,7 @@ func Test_Branch_Copy(t *testing.T) {
 
 	testCases := map[string]struct {
 		branch         *Branch
-		copyChildren   bool
+		settings       CopySettings
 		expectedBranch *Branch
 	}{
 		"empty branch": {
@@ -44,15 +44,14 @@ func Test_Branch_Copy(t *testing.T) {
 				HashDigest: []byte{5},
 				Encoding:   []byte{6},
 			},
+			settings: DefaultCopySettings,
 			expectedBranch: &Branch{
 				Key:   []byte{1, 2},
 				Value: []byte{3, 4},
 				Children: [16]Node{
 					nil, nil, &Leaf{Key: []byte{9}},
 				},
-				Dirty:      true,
-				HashDigest: []byte{5},
-				Encoding:   []byte{6},
+				Dirty: true,
 			},
 		},
 		"branch with children copied": {
@@ -61,21 +60,46 @@ func Test_Branch_Copy(t *testing.T) {
 					nil, nil, &Leaf{Key: []byte{9}},
 				},
 			},
-			copyChildren: true,
+			settings: CopySettings{
+				CopyChildren: true,
+			},
 			expectedBranch: &Branch{
 				Children: [16]Node{
 					nil, nil, &Leaf{Key: []byte{9}},
 				},
 			},
 		},
+		"deep copy": {
+			branch: &Branch{
+				Key:   []byte{1, 2},
+				Value: []byte{3, 4},
+				Children: [16]Node{
+					nil, nil, &Leaf{Key: []byte{9}},
+				},
+				Dirty:      true,
+				HashDigest: []byte{5},
+				Encoding:   []byte{6},
+			},
+			settings: DeepCopySettings,
+			expectedBranch: &Branch{
+				Key:   []byte{1, 2},
+				Value: []byte{3, 4},
+				Children: [16]Node{
+					nil, nil, &Leaf{Key: []byte{9}},
+				},
+				Dirty:      true,
+				HashDigest: []byte{5},
+				Encoding:   []byte{6},
+			},
+		},
 	}
 
 	for name, testCase := range testCases {
 		testCase := testCase
 		t.Run(name, func(t *testing.T) {
 			t.Parallel()
 
-			nodeCopy := testCase.branch.Copy(testCase.copyChildren)
+			nodeCopy := testCase.branch.Copy(testCase.settings)
 
 			branchCopy, ok := nodeCopy.(*Branch)
 			require.True(t, ok)
@@ -97,10 +121,12 @@ func Test_Leaf_Copy(t *testing.T) {
 
 	testCases := map[string]struct {
 		leaf         *Leaf
+		settings     CopySettings
 		expectedLeaf *Leaf
 	}{
 		"empty leaf": {
 			leaf:         &Leaf{},
+			settings:     DefaultCopySettings,
 			expectedLeaf: &Leaf{},
 		},
 		"non empty leaf": {
@@ -111,6 +137,22 @@ func Test_Leaf_Copy(t *testing.T) {
 				HashDigest: []byte{5},
 				Encoding:   []byte{6},
 			},
+			settings: DefaultCopySettings,
+			expectedLeaf: &Leaf{
+				Key:   []byte{1, 2},
+				Value: []byte{3, 4},
+				Dirty: true,
+			},
+		},
+		"deep copy": {
+			leaf: &Leaf{
+				Key:        []byte{1, 2},
+				Value:      []byte{3, 4},
+				Dirty:      true,
+				HashDigest: []byte{5},
+				Encoding:   []byte{6},
+			},
+			settings: DeepCopySettings,
 			expectedLeaf: &Leaf{
 				Key:        []byte{1, 2},
 				Value:      []byte{3, 4},
@@ -126,8 +168,7 @@ func Test_Leaf_Copy(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			t.Parallel()
 
-			const copyChildren = false
-			nodeCopy := testCase.leaf.Copy(copyChildren)
+			nodeCopy := testCase.leaf.Copy(testCase.settings)
 
 			leafCopy, ok := nodeCopy.(*Leaf)
 			require.True(t, ok)

diff --git a/internal/trie/node/node.go b/internal/trie/node/node.go
@@ -20,7 +20,7 @@ type Node interface {
 	GetValue() (value []byte)
 	GetGeneration() (generation uint64)
 	SetGeneration(generation uint64)
-	Copy(copyChildren bool) Node
+	Copy(settings CopySettings) (cpy Node)
 	Type() Type
 	StringNode() (stringNode *gotree.Node)
 }