chore(lib/trie): minimal refactor for more testable trie production c…

…ode (#2144)

- Detach methods from `Trie` receiver to simple functions: `getKeysWithPrefix`, `addAllKeys`, `nextKey`, `entries`, `encodeRoot` 
- `encodeRoot` uses `node.Buffer` interface
- Add some minor comments
- Fix/improve `Trie.DeepCopy()` used in tests
- Change Copy method to copy children optionally (for trie DeepCopy)
- (Re-)export `Encoding` field in branches and leaves node
- (Re-)export `Dirty` field in branches and leaves node
- Simplify `deleteNodes`

dot/rpc/modules/author_integration_test.go

@@ -176,7 +176,7 @@ func TestAuthorModule_SubmitExtrinsic_invalid_input(t *testing.T) {
 
 	res := new(ExtrinsicHashResponse)
 	err := auth.SubmitExtrinsic(nil, &ext, res)
-	require.EqualError(t, err, "could not byteify non 0x prefixed string")
+	require.EqualError(t, err, "could not byteify non 0x prefixed string: 0x31")
 }
 
 func TestAuthorModule_SubmitExtrinsic_InQueue(t *testing.T) {

dot/rpc/modules/author_test.go

@@ -78,7 +78,7 @@ func TestAuthorModule_HasSessionKeys(t *testing.T) {
 				req: &HasSessionKeyRequest{},
 			},
 			exp:    false,
-			expErr: errors.New("invalid string"),
+			expErr: errors.New("could not byteify non 0x prefixed string: "),
 		},
 		{
 			name: "decodeSessionKeys err",
@@ -179,7 +179,7 @@ func TestAuthorModule_SubmitExtrinsic(t *testing.T) {
 			args: args{
 				req: &Extrinsic{fmt.Sprintf("%x", "1")},
 			},
-			expErr:  fmt.Errorf("could not byteify non 0x prefixed string"),
+			expErr:  fmt.Errorf("could not byteify non 0x prefixed string: 31"),
 			wantRes: ExtrinsicHashResponse(""),
 		},
 		{

dot/rpc/modules/offchain_test.go

@@ -86,7 +86,7 @@ func TestOffchainModule_LocalStorageGet(t *testing.T) {
 					Key:  "0x1",
 				},
 			},
-			expErr: errors.New("cannot decode an odd length string"),
+			expErr: errors.New("encoding/hex: odd length hex string: 0x1"),
 		},
 	}
 	for _, tt := range tests {
@@ -153,7 +153,7 @@ func TestOffchainModule_LocalStorageSet(t *testing.T) {
 					Value: "0x22222222222222",
 				},
 			},
-			expErr: errors.New("cannot decode an odd length string"),
+			expErr: errors.New("encoding/hex: odd length hex string: 0x1"),
 		},
 		{
 			name: "Invalid Value",
@@ -167,7 +167,7 @@ func TestOffchainModule_LocalStorageSet(t *testing.T) {
 					Value: "0x2",
 				},
 			},
-			expErr: errors.New("cannot decode an odd length string"),
+			expErr: errors.New("encoding/hex: odd length hex string: 0x2"),
 		},
 		{
 			name: "setPersistentError",

dot/rpc/modules/payment_test.go

@@ -108,7 +108,7 @@ func TestPaymentModule_QueryInfo(t *testing.T) {
 					Ext: "0x0",
 				},
 			},
-			expErr: errors.New("cannot decode an odd length string"),
+			expErr: errors.New("encoding/hex: odd length hex string: 0x0"),
 		},
 		{
 			name: "Invalid Ext",

dot/rpc/modules/state_test.go

@@ -256,7 +256,7 @@ func TestStateModuleGetKeysPaged(t *testing.T) {
 					AfterKey: "0x01",
 				},
 			},
-			expErr: errors.New("invalid string"),
+			expErr: errors.New("could not byteify non 0x prefixed string: a"),
 		},
 	}
 	for _, tt := range tests {

dot/rpc/subscription/websocket_test.go

@@ -209,7 +209,7 @@ func TestWSConn_HandleComm(t *testing.T) {
 	wsconn.BlockAPI = nil
 	wsconn.TxStateAPI = modules.NewMockTransactionStateAPI()
 	listner, err := wsconn.initExtrinsicWatch(0, []interface{}{"NotHex"})
-	require.EqualError(t, err, "could not byteify non 0x prefixed string")
+	require.EqualError(t, err, "could not byteify non 0x prefixed string: NotHex")
 	require.Nil(t, listner)
 
 	listner, err = wsconn.initExtrinsicWatch(0, []interface{}{"0x26aa"})

dot/state/service_test.go

@@ -277,12 +277,10 @@ func TestService_PruneStorage(t *testing.T) {
 		require.NoError(t, err)
 
 		// Store the other blocks that will be pruned.
-		var trieVal *trie.Trie
-		trieVal, err = trieState.Trie().DeepCopy()
-		require.NoError(t, err)
+		copiedTrie := trieState.Trie().DeepCopy()
 
 		var rootHash common.Hash
-		rootHash, err = trieVal.Hash()
+		rootHash, err = copiedTrie.Hash()
 		require.NoError(t, err)
 
 		prunedArr = append(prunedArr, prunedBlock{hash: block.Header.StateRoot, dbKey: rootHash[:]})

dot/utils_test.go

@@ -69,8 +69,7 @@ func TestTrieSnapshot(t *testing.T) {
 		tri.Put([]byte(k), val)
 	}
 
-	// DeepCopy the trie.
-	dcTrie, err := tri.DeepCopy()
+	dcTrie := tri.DeepCopy()
 	require.NoError(t, err)
 
 	// Take Snapshot of the trie.

internal/trie/node/branch.go

@@ -14,19 +14,20 @@ var _ Node = (*Branch)(nil)
 
 // Branch is a branch in the trie.
 type Branch struct {
-	Key      []byte // partial key
+	// Partial key bytes in nibbles (0 to f in hexadecimal)
+	Key      []byte
 	Children [16]Node
 	Value    []byte
-	// dirty is true when the branch differs
+	// Dirty is true when the branch differs
 	// from the node stored in the database.
-	dirty      bool
+	Dirty      bool
 	hashDigest []byte
-	encoding   []byte
-	// generation is incremented on every trie Snapshot() call.
-	// Each node also contain a certain generation number,
-	// which is updated to match the trie generation once they are
+	Encoding   []byte
+	// Generation is incremented on every trie Snapshot() call.
+	// Each node also contain a certain Generation number,
+	// which is updated to match the trie Generation once they are
 	// inserted, moved or iterated over.
-	generation uint64
+	Generation uint64
 	sync.RWMutex
 }
 
@@ -35,8 +36,8 @@ func NewBranch(key, value []byte, dirty bool, generation uint64) *Branch {
 	return &Branch{
 		Key:        key,
 		Value:      value,
-		dirty:      dirty,
-		generation: generation,
+		Dirty:      dirty,
+		Generation: generation,
 	}
 }
 
@@ -52,8 +53,8 @@ func (b *Branch) Type() Type {
 func (b *Branch) String() string {
 	if len(b.Value) > 1024 {
 		return fmt.Sprintf("branch key=0x%x childrenBitmap=%b value (hashed)=0x%x dirty=%t",
-			b.Key, b.ChildrenBitmap(), common.MustBlake2bHash(b.Value), b.dirty)
+			b.Key, b.ChildrenBitmap(), common.MustBlake2bHash(b.Value), b.Dirty)
 	}
 	return fmt.Sprintf("branch key=0x%x childrenBitmap=%b value=0x%x dirty=%t",
-		b.Key, b.ChildrenBitmap(), b.Value, b.dirty)
+		b.Key, b.ChildrenBitmap(), b.Value, b.Dirty)
 }

internal/trie/node/branch_encode.go

@@ -77,8 +77,8 @@ func (b *Branch) hash(digestBuffer io.Writer) (err error) {
 // Encode encodes a branch with the encoding specified at the top of this package
 // to the buffer given.
 func (b *Branch) Encode(buffer Buffer) (err error) {
-	if !b.dirty && b.encoding != nil {
-		_, err = buffer.Write(b.encoding)
+	if !b.Dirty && b.Encoding != nil {
+		_, err = buffer.Write(b.Encoding)
 		if err != nil {
 			return fmt.Errorf("cannot write stored encoding to buffer: %w", err)
 		}

internal/trie/node/branch_encode_test.go

@@ -152,7 +152,7 @@ func Test_Branch_Encode(t *testing.T) {
 	}{
 		"clean branch with encoding": {
 			branch: &Branch{
-				encoding: []byte{1, 2, 3},
+				Encoding: []byte{1, 2, 3},
 			},
 			writes: []writeCall{
 				{ // stored encoding
@@ -162,7 +162,7 @@ func Test_Branch_Encode(t *testing.T) {
 		},
 		"write error for clean branch with encoding": {
 			branch: &Branch{
-				encoding: []byte{1, 2, 3},
+				Encoding: []byte{1, 2, 3},
 			},
 			writes: []writeCall{
 				{ // stored encoding

internal/trie/node/branch_test.go

@@ -22,8 +22,8 @@ func Test_NewBranch(t *testing.T) {
 	expectedBranch := &Branch{
 		Key:        key,
 		Value:      value,
-		dirty:      dirty,
-		generation: generation,
+		Dirty:      dirty,
+		Generation: generation,
 	}
 	assert.Equal(t, expectedBranch, branch)
 
@@ -83,7 +83,7 @@ func Test_Branch_String(t *testing.T) {
 			branch: &Branch{
 				Key:   []byte{1, 2},
 				Value: []byte{3, 4},
-				dirty: true,
+				Dirty: true,
 				Children: [16]Node{
 					nil, nil, nil,
 					&Leaf{},
@@ -100,7 +100,7 @@ func Test_Branch_String(t *testing.T) {
 			branch: &Branch{
 				Key:   []byte{1, 2},
 				Value: make([]byte, 1025),
-				dirty: true,
+				Dirty: true,
 				Children: [16]Node{
 					nil, nil, nil,
 					&Leaf{},

internal/trie/node/copy.go

@@ -4,16 +4,27 @@
 package node
 
 // Copy deep copies the branch.
-func (b *Branch) Copy() Node {
+// Setting copyChildren to true will deep copy
+// children as well.
+func (b *Branch) Copy(copyChildren bool) Node {
 	b.RLock()
 	defer b.RUnlock()
 
 	cpy := &Branch{
-		Children:   b.Children, // copy interface pointers
-		dirty:      b.dirty,
-		generation: b.generation,
+		Dirty:      b.Dirty,
+		Generation: b.Generation,
+	}
+
+	if copyChildren {
+		for i, child := range b.Children {
+			if child == nil {
+				continue
+			}
+			cpy.Children[i] = child.Copy(copyChildren)
+		}
+	} else {
+		cpy.Children = b.Children // copy interface pointers only
 	}
-	copy(cpy.Key, b.Key)
 
 	if b.Key != nil {
 		cpy.Key = make([]byte, len(b.Key))
@@ -31,25 +42,25 @@ func (b *Branch) Copy() Node {
 		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() Node {
+func (l *Leaf) Copy(_ bool) Node {
 	l.RLock()
 	defer l.RUnlock()
 
 	l.encodingMu.RLock()
 	defer l.encodingMu.RUnlock()
 
 	cpy := &Leaf{
-		dirty:      l.dirty,
-		generation: l.generation,
+		Dirty:      l.Dirty,
+		Generation: l.Generation,
 	}
 
 	if l.Key != nil {
@@ -68,9 +79,9 @@ func (l *Leaf) Copy() Node {
 		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

internal/trie/node/copy_test.go

@@ -26,6 +26,7 @@ func Test_Branch_Copy(t *testing.T) {
 
 	testCases := map[string]struct {
 		branch         *Branch
+		copyChildren   bool
 		expectedBranch *Branch
 	}{
 		"empty branch": {
@@ -39,19 +40,32 @@ func Test_Branch_Copy(t *testing.T) {
 				Children: [16]Node{
 					nil, nil, &Leaf{Key: []byte{9}},
 				},
-				dirty:      true,
+				Dirty:      true,
 				hashDigest: []byte{5},
-				encoding:   []byte{6},
+				Encoding:   []byte{6},
 			},
 			expectedBranch: &Branch{
 				Key:   []byte{1, 2},
 				Value: []byte{3, 4},
 				Children: [16]Node{
 					nil, nil, &Leaf{Key: []byte{9}},
 				},
-				dirty:      true,
+				Dirty:      true,
 				hashDigest: []byte{5},
-				encoding:   []byte{6},
+				Encoding:   []byte{6},
+			},
+		},
+		"branch with children copied": {
+			branch: &Branch{
+				Children: [16]Node{
+					nil, nil, &Leaf{Key: []byte{9}},
+				},
+			},
+			copyChildren: true,
+			expectedBranch: &Branch{
+				Children: [16]Node{
+					nil, nil, &Leaf{Key: []byte{9}},
+				},
 			},
 		},
 	}
@@ -61,7 +75,7 @@ func Test_Branch_Copy(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			t.Parallel()
 
-			nodeCopy := testCase.branch.Copy()
+			nodeCopy := testCase.branch.Copy(testCase.copyChildren)
 
 			branchCopy, ok := nodeCopy.(*Branch)
 			require.True(t, ok)
@@ -70,7 +84,7 @@ func Test_Branch_Copy(t *testing.T) {
 			testForSliceModif(t, testCase.branch.Key, branchCopy.Key)
 			testForSliceModif(t, testCase.branch.Value, branchCopy.Value)
 			testForSliceModif(t, testCase.branch.hashDigest, branchCopy.hashDigest)
-			testForSliceModif(t, testCase.branch.encoding, branchCopy.encoding)
+			testForSliceModif(t, testCase.branch.Encoding, branchCopy.Encoding)
 
 			testCase.branch.Children[15] = &Leaf{Key: []byte("modified")}
 			assert.NotEqual(t, branchCopy.Children, testCase.branch.Children)
@@ -93,16 +107,16 @@ func Test_Leaf_Copy(t *testing.T) {
 			leaf: &Leaf{
 				Key:        []byte{1, 2},
 				Value:      []byte{3, 4},
-				dirty:      true,
+				Dirty:      true,
 				hashDigest: []byte{5},
-				encoding:   []byte{6},
+				Encoding:   []byte{6},
 			},
 			expectedLeaf: &Leaf{
 				Key:        []byte{1, 2},
 				Value:      []byte{3, 4},
-				dirty:      true,
+				Dirty:      true,
 				hashDigest: []byte{5},
-				encoding:   []byte{6},
+				Encoding:   []byte{6},
 			},
 		},
 	}
@@ -112,7 +126,8 @@ func Test_Leaf_Copy(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			t.Parallel()
 
-			nodeCopy := testCase.leaf.Copy()
+			const copyChildren = false
+			nodeCopy := testCase.leaf.Copy(copyChildren)
 
 			leafCopy, ok := nodeCopy.(*Leaf)
 			require.True(t, ok)
@@ -121,7 +136,7 @@ func Test_Leaf_Copy(t *testing.T) {
 			testForSliceModif(t, testCase.leaf.Key, leafCopy.Key)
 			testForSliceModif(t, testCase.leaf.Value, leafCopy.Value)
 			testForSliceModif(t, testCase.leaf.hashDigest, leafCopy.hashDigest)
-			testForSliceModif(t, testCase.leaf.encoding, leafCopy.encoding)
+			testForSliceModif(t, testCase.leaf.Encoding, leafCopy.Encoding)
 		})
 	}
 }