Trie code tracks number of descendants for each node

- `clearPrefix`: use `nodesRemoved == 0` instead of `updated` bool
- Add more test cases
- Remove old irrelevant `TODO`
- Print descendants in node String methods
- Add descendants check fuzz test

internal/trie/node/branch.go

@@ -30,6 +30,12 @@ type Branch struct {
 	// which is updated to match the trie Generation once they are
 	// inserted, moved or iterated over.
 	Generation uint64
+
+	// Statistics
+
+	// Descendants is the number of descendant nodes for
+	// this particular node.
+	Descendants uint32
 }
 
 // NewBranch creates a new branch using the arguments given.
@@ -62,6 +68,7 @@ func (b *Branch) StringNode() (stringNode *gotree.Node) {
 	stringNode.Appendf("Dirty: %t", b.Dirty)
 	stringNode.Appendf("Key: " + bytesToString(b.Key))
 	stringNode.Appendf("Value: " + bytesToString(b.Value))
+	stringNode.Appendf("Descendants: %d", b.Descendants)
 	stringNode.Appendf("Calculated encoding: " + bytesToString(b.Encoding))
 	stringNode.Appendf("Calculated digest: " + bytesToString(b.HashDigest))
 

internal/trie/node/branch_test.go

@@ -82,14 +82,16 @@ func Test_Branch_String(t *testing.T) {
 ├── Dirty: false
 ├── Key: nil
 ├── Value: nil
+├── Descendants: 0
 ├── Calculated encoding: nil
 └── Calculated digest: nil`,
 		},
 		"branch with value smaller than 1024": {
 			branch: &Branch{
-				Key:   []byte{1, 2},
-				Value: []byte{3, 4},
-				Dirty: true,
+				Key:         []byte{1, 2},
+				Value:       []byte{3, 4},
+				Dirty:       true,
+				Descendants: 3,
 				Children: [16]Node{
 					nil, nil, nil,
 					&Leaf{},
@@ -105,6 +107,7 @@ func Test_Branch_String(t *testing.T) {
 ├── Dirty: true
 ├── Key: 0x0102
 ├── Value: 0x0304
+├── Descendants: 3
 ├── Calculated encoding: nil
 ├── Calculated digest: nil
 ├── Child 3
@@ -121,6 +124,7 @@ func Test_Branch_String(t *testing.T) {
 |       ├── Dirty: false
 |       ├── Key: nil
 |       ├── Value: nil
+|       ├── Descendants: 0
 |       ├── Calculated encoding: nil
 |       └── Calculated digest: nil
 └── Child 11
@@ -134,9 +138,10 @@ func Test_Branch_String(t *testing.T) {
 		},
 		"branch with value higher than 1024": {
 			branch: &Branch{
-				Key:   []byte{1, 2},
-				Value: make([]byte, 1025),
-				Dirty: true,
+				Key:         []byte{1, 2},
+				Value:       make([]byte, 1025),
+				Dirty:       true,
+				Descendants: 3,
 				Children: [16]Node{
 					nil, nil, nil,
 					&Leaf{},
@@ -152,6 +157,7 @@ func Test_Branch_String(t *testing.T) {
 ├── Dirty: true
 ├── Key: 0x0102
 ├── Value: 0x0000000000000000...0000000000000000
+├── Descendants: 3
 ├── Calculated encoding: nil
 ├── Calculated digest: nil
 ├── Child 3
@@ -168,6 +174,7 @@ func Test_Branch_String(t *testing.T) {
 |       ├── Dirty: false
 |       ├── Key: nil
 |       ├── Value: nil
+|       ├── Descendants: 0
 |       ├── Calculated encoding: nil
 |       └── Calculated digest: nil
 └── Child 11

internal/trie/node/copy.go

@@ -57,8 +57,9 @@ type CopySettings struct {
 // children as well.
 func (b *Branch) Copy(settings CopySettings) Node {
 	cpy := &Branch{
-		Dirty:      b.Dirty,
-		Generation: b.Generation,
+		Dirty:       b.Dirty,
+		Generation:  b.Generation,
+		Descendants: b.GetDescendants(),
 	}
 
 	if settings.CopyChildren {

internal/trie/node/decode.go

@@ -98,6 +98,7 @@ func decodeBranch(reader io.Reader, header byte) (branch *Branch, err error) {
 		if (childrenBitmap[i/8]>>(i%8))&1 != 1 {
 			continue
 		}
+
 		var hash []byte
 		err := sd.Decode(&hash)
 		if err != nil {
@@ -113,10 +114,12 @@ func decodeBranch(reader io.Reader, header byte) (branch *Branch, err error) {
 				return nil, fmt.Errorf("%w: at index %d: %s",
 					ErrDecodeValue, i, err)
 			}
+			branch.AddDescendants(1)
 			branch.Children[i] = leaf
 			continue
 		}
 
+		branch.AddDescendants(1)
 		branch.Children[i] = &Leaf{
 			HashDigest: hash,
 		}

internal/trie/node/decode_test.go

@@ -119,6 +119,7 @@ func Test_Decode(t *testing.T) {
 					13, 10, 0, 14, 7, 10,
 					4, 1, 1, 3, 12, 4,
 				},
+				Descendants: 2,
 				Children: [16]Node{
 					nil, nil, nil, nil,
 					&Leaf{
@@ -235,7 +236,8 @@ func Test_decodeBranch(t *testing.T) {
 						HashDigest: []byte{1, 2, 3, 4, 5},
 					},
 				},
-				Dirty: true,
+				Dirty:       true,
+				Descendants: 1,
 			},
 		},
 		"value decoding error for node type 3": {
@@ -270,7 +272,8 @@ func Test_decodeBranch(t *testing.T) {
 						HashDigest: []byte{1, 2, 3, 4, 5},
 					},
 				},
-				Dirty: true,
+				Dirty:       true,
+				Descendants: 1,
 			},
 		},
 	}

internal/trie/node/encode_decode_test.go

@@ -54,7 +54,8 @@ func Test_Branch_Encode_Decode(t *testing.T) {
 				},
 			},
 			branchDecoded: &Branch{
-				Key: []byte{5},
+				Key:         []byte{5},
+				Descendants: 1,
 				Children: [16]Node{
 					&Leaf{
 						Key:   []byte{9},
@@ -103,7 +104,8 @@ func Test_Branch_Encode_Decode(t *testing.T) {
 						},
 					},
 				},
-				Dirty: true,
+				Dirty:       true,
+				Descendants: 1,
 			},
 		},
 	}

internal/trie/node/stats.go

@@ -0,0 +1,19 @@
+// Copyright 2022 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+package node
+
+// GetDescendants returns the number of descendants in the branch.
+func (b *Branch) GetDescendants() (descendants uint32) {
+	return b.Descendants
+}
+
+// AddDescendants adds descendant nodes count to the node stats.
+func (b *Branch) AddDescendants(n uint32) {
+	b.Descendants += n
+}
+
+// SubDescendants subtracts descendant nodes count from the node stats.
+func (b *Branch) SubDescendants(n uint32) {
+	b.Descendants -= n
+}

internal/trie/node/stats_test.go

@@ -0,0 +1,60 @@
+// Copyright 2022 ChainSafe Systems (ON)
+// SPDX-License-Identifier: LGPL-3.0-only
+
+package node
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_Branch_GetDescendants(t *testing.T) {
+	t.Parallel()
+
+	const descendants uint32 = 10
+	branch := &Branch{
+		Descendants: descendants,
+	}
+	result := branch.GetDescendants()
+
+	assert.Equal(t, descendants, result)
+}
+
+func Test_Branch_AddDescendants(t *testing.T) {
+	t.Parallel()
+
+	const (
+		initialDescendants uint32 = 10
+		addDescendants     uint32 = 2
+		finalDescendants   uint32 = 12
+	)
+	branch := &Branch{
+		Descendants: initialDescendants,
+	}
+	branch.AddDescendants(addDescendants)
+	expected := &Branch{
+		Descendants: finalDescendants,
+	}
+
+	assert.Equal(t, expected, branch)
+}
+
+func Test_Branch_SubDescendants(t *testing.T) {
+	t.Parallel()
+
+	const (
+		initialDescendants uint32 = 10
+		subDescendants     uint32 = 2
+		finalDescendants   uint32 = 8
+	)
+	branch := &Branch{
+		Descendants: initialDescendants,
+	}
+	branch.SubDescendants(subDescendants)
+	expected := &Branch{
+		Descendants: finalDescendants,
+	}
+
+	assert.Equal(t, expected, branch)
+}

lib/trie/database.go

@@ -219,6 +219,17 @@ func (t *Trie) load(db chaindb.Database, n Node) error {
 		if err != nil {
 			return fmt.Errorf("cannot load child at index %d with hash 0x%x: %w", i, hash, err)
 		}
+
+		if decodedNode.Type() != node.LeafType { // branch decoded node
+			// Note 1: the node is fully loaded with all its descendants
+			// count only after the database load above.
+			// Note 2: direct child node is already counted as descendant
+			// when it was read as a leaf with hash only in decodeBranch,
+			// so we only add the descendants of the child branch to the
+			// current branch.
+			childBranchDescendants := decodedNode.(*node.Branch).Descendants
+			branch.AddDescendants(childBranchDescendants)
+		}
 	}
 
 	for _, key := range t.GetKeysWithPrefix(ChildStorageKeyPrefix) {

lib/trie/print_test.go

@@ -39,8 +39,9 @@ func Test_Trie_String(t *testing.T) {
 		"branch root": {
 			trie: Trie{
 				root: &node.Branch{
-					Key:   nil,
-					Value: []byte{1, 2},
+					Key:         nil,
+					Value:       []byte{1, 2},
+					Descendants: 2,
 					Children: [16]node.Node{
 						&node.Leaf{
 							Key:        []byte{1, 2, 3},
@@ -61,6 +62,7 @@ func Test_Trie_String(t *testing.T) {
 ├── Dirty: false
 ├── Key: nil
 ├── Value: 0x0102
+├── Descendants: 2
 ├── Calculated encoding: nil
 ├── Calculated digest: nil
 ├── Child 0