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

* Note: descendants are stored in every node since we have a delete operation that cuts of an entire branch, so it's more performant to have it get the count of descendants to remove from the branch being removed, rather than iterating over this branch and all its children.
* Adapt existing tests
* Add test cases to reach back full test coverage

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