feat(lib/trie): add descendants count to branches

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

@@ -8,8 +8,9 @@ package node
 // children as well.
 func (b *Branch) Copy(copyChildren bool) Node {
 	cpy := &Branch{
-		Dirty:      b.Dirty,
-		Generation: b.Generation,
+		Dirty:       b.Dirty,
+		Generation:  b.Generation,
+		Descendants: b.GetDescendants(),
 	}
 
 	if copyChildren {

internal/trie/node/decode.go

@@ -98,6 +98,8 @@ func decodeBranch(reader io.Reader, header byte) (branch *Branch, err error) {
 		if (childrenBitmap[i/8]>>(i%8))&1 != 1 {
 			continue
 		}
+		branch.AddDescendants(1)
+
 		var hash []byte
 		err := sd.Decode(&hash)
 		if err != nil {

internal/trie/node/decode_test.go

@@ -176,7 +176,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": {
@@ -211,7 +212,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

@@ -60,7 +60,8 @@ func Test_Branch_Encode_Decode(t *testing.T) {
 						HashDigest: []byte{0x41, 0x9, 0x4, 0xa},
 					},
 				},
-				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/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