Skip to content

Commit 6b04e39

Browse files
qdm12qdm12
committed
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
1 parent 91e7ffa commit 6b04e39

File tree

13 files changed

+977
-398
lines changed

13 files changed

+977
-398
lines changed

‎internal/trie/node/branch.go

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,12 @@ type Branch struct {
3030
// which is updated to match the trie Generation once they are
3131
// inserted, moved or iterated over.
3232
Generation uint64
33+
34+
// Statistics
35+
36+
// Descendants is the number of descendant nodes for
37+
// this particular node.
38+
Descendants uint32
3339
}
3440

3541
// NewBranch creates a new branch using the arguments given.
@@ -62,6 +68,7 @@ func (b *Branch) StringNode() (stringNode *gotree.Node) {
6268
stringNode.Appendf("Dirty: %t", b.Dirty)
6369
stringNode.Appendf("Key: " + bytesToString(b.Key))
6470
stringNode.Appendf("Value: " + bytesToString(b.Value))
71+
stringNode.Appendf("Descendants: %d", b.Descendants)
6572
stringNode.Appendf("Calculated encoding: " + bytesToString(b.Encoding))
6673
stringNode.Appendf("Calculated digest: " + bytesToString(b.HashDigest))
6774

‎internal/trie/node/branch_test.go

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -82,14 +82,16 @@ func Test_Branch_String(t *testing.T) {
8282
├── Dirty: false
8383
├── Key: nil
8484
├── Value: nil
85+
├── Descendants: 0
8586
├── Calculated encoding: nil
8687
└── Calculated digest: nil`,
8788
},
8889
"branch with value smaller than 1024": {
8990
branch: &Branch{
90-
Key: []byte{1, 2},
91-
Value: []byte{3, 4},
92-
Dirty: true,
91+
Key: []byte{1, 2},
92+
Value: []byte{3, 4},
93+
Dirty: true,
94+
Descendants: 3,
9395
Children: [16]Node{
9496
nil, nil, nil,
9597
&Leaf{},
@@ -105,6 +107,7 @@ func Test_Branch_String(t *testing.T) {
105107
├── Dirty: true
106108
├── Key: 0x0102
107109
├── Value: 0x0304
110+
├── Descendants: 3
108111
├── Calculated encoding: nil
109112
├── Calculated digest: nil
110113
├── Child 3
@@ -121,6 +124,7 @@ func Test_Branch_String(t *testing.T) {
121124
| ├── Dirty: false
122125
| ├── Key: nil
123126
| ├── Value: nil
127+
| ├── Descendants: 0
124128
| ├── Calculated encoding: nil
125129
| └── Calculated digest: nil
126130
└── Child 11
@@ -134,9 +138,10 @@ func Test_Branch_String(t *testing.T) {
134138
},
135139
"branch with value higher than 1024": {
136140
branch: &Branch{
137-
Key: []byte{1, 2},
138-
Value: make([]byte, 1025),
139-
Dirty: true,
141+
Key: []byte{1, 2},
142+
Value: make([]byte, 1025),
143+
Dirty: true,
144+
Descendants: 3,
140145
Children: [16]Node{
141146
nil, nil, nil,
142147
&Leaf{},
@@ -152,6 +157,7 @@ func Test_Branch_String(t *testing.T) {
152157
├── Dirty: true
153158
├── Key: 0x0102
154159
├── Value: 0x0000000000000000...0000000000000000
160+
├── Descendants: 3
155161
├── Calculated encoding: nil
156162
├── Calculated digest: nil
157163
├── Child 3
@@ -168,6 +174,7 @@ func Test_Branch_String(t *testing.T) {
168174
| ├── Dirty: false
169175
| ├── Key: nil
170176
| ├── Value: nil
177+
| ├── Descendants: 0
171178
| ├── Calculated encoding: nil
172179
| └── Calculated digest: nil
173180
└── Child 11

‎internal/trie/node/copy.go

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -57,8 +57,9 @@ type CopySettings struct {
5757
// children as well.
5858
func (b *Branch) Copy(settings CopySettings) Node {
5959
cpy := &Branch{
60-
Dirty: b.Dirty,
61-
Generation: b.Generation,
60+
Dirty: b.Dirty,
61+
Generation: b.Generation,
62+
Descendants: b.GetDescendants(),
6263
}
6364

6465
if settings.CopyChildren {

‎internal/trie/node/decode.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@ func decodeBranch(reader io.Reader, header byte) (branch *Branch, err error) {
9898
if (childrenBitmap[i/8]>>(i%8))&1 != 1 {
9999
continue
100100
}
101+
101102
var hash []byte
102103
err := sd.Decode(&hash)
103104
if err != nil {
@@ -113,10 +114,12 @@ func decodeBranch(reader io.Reader, header byte) (branch *Branch, err error) {
113114
return nil, fmt.Errorf("%w: at index %d: %s",
114115
ErrDecodeValue, i, err)
115116
}
117+
branch.AddDescendants(1)
116118
branch.Children[i] = leaf
117119
continue
118120
}
119121

122+
branch.AddDescendants(1)
120123
branch.Children[i] = &Leaf{
121124
HashDigest: hash,
122125
}

‎internal/trie/node/decode_test.go

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,7 @@ func Test_Decode(t *testing.T) {
119119
13, 10, 0, 14, 7, 10,
120120
4, 1, 1, 3, 12, 4,
121121
},
122+
Descendants: 2,
122123
Children: [16]Node{
123124
nil, nil, nil, nil,
124125
&Leaf{
@@ -235,7 +236,8 @@ func Test_decodeBranch(t *testing.T) {
235236
HashDigest: []byte{1, 2, 3, 4, 5},
236237
},
237238
},
238-
Dirty: true,
239+
Dirty: true,
240+
Descendants: 1,
239241
},
240242
},
241243
"value decoding error for node type 3": {
@@ -270,7 +272,8 @@ func Test_decodeBranch(t *testing.T) {
270272
HashDigest: []byte{1, 2, 3, 4, 5},
271273
},
272274
},
273-
Dirty: true,
275+
Dirty: true,
276+
Descendants: 1,
274277
},
275278
},
276279
}

‎internal/trie/node/encode_decode_test.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,8 @@ func Test_Branch_Encode_Decode(t *testing.T) {
5454
},
5555
},
5656
branchDecoded: &Branch{
57-
Key: []byte{5},
57+
Key: []byte{5},
58+
Descendants: 1,
5859
Children: [16]Node{
5960
&Leaf{
6061
Key: []byte{9},
@@ -103,7 +104,8 @@ func Test_Branch_Encode_Decode(t *testing.T) {
103104
},
104105
},
105106
},
106-
Dirty: true,
107+
Dirty: true,
108+
Descendants: 1,
107109
},
108110
},
109111
}

‎internal/trie/node/stats.go

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
// Copyright 2022 ChainSafe Systems (ON)
2+
// SPDX-License-Identifier: LGPL-3.0-only
3+
4+
package node
5+
6+
// GetDescendants returns the number of descendants in the branch.
7+
func (b *Branch) GetDescendants() (descendants uint32) {
8+
return b.Descendants
9+
}
10+
11+
// AddDescendants adds descendant nodes count to the node stats.
12+
func (b *Branch) AddDescendants(n uint32) {
13+
b.Descendants += n
14+
}
15+
16+
// SubDescendants subtracts descendant nodes count from the node stats.
17+
func (b *Branch) SubDescendants(n uint32) {
18+
b.Descendants -= n
19+
}

‎internal/trie/node/stats_test.go

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
// Copyright 2022 ChainSafe Systems (ON)
2+
// SPDX-License-Identifier: LGPL-3.0-only
3+
4+
package node
5+
6+
import (
7+
"testing"
8+
9+
"github.com/stretchr/testify/assert"
10+
)
11+
12+
func Test_Branch_GetDescendants(t *testing.T) {
13+
t.Parallel()
14+
15+
const descendants uint32 = 10
16+
branch := &Branch{
17+
Descendants: descendants,
18+
}
19+
result := branch.GetDescendants()
20+
21+
assert.Equal(t, descendants, result)
22+
}
23+
24+
func Test_Branch_AddDescendants(t *testing.T) {
25+
t.Parallel()
26+
27+
const (
28+
initialDescendants uint32 = 10
29+
addDescendants uint32 = 2
30+
finalDescendants uint32 = 12
31+
)
32+
branch := &Branch{
33+
Descendants: initialDescendants,
34+
}
35+
branch.AddDescendants(addDescendants)
36+
expected := &Branch{
37+
Descendants: finalDescendants,
38+
}
39+
40+
assert.Equal(t, expected, branch)
41+
}
42+
43+
func Test_Branch_SubDescendants(t *testing.T) {
44+
t.Parallel()
45+
46+
const (
47+
initialDescendants uint32 = 10
48+
subDescendants uint32 = 2
49+
finalDescendants uint32 = 8
50+
)
51+
branch := &Branch{
52+
Descendants: initialDescendants,
53+
}
54+
branch.SubDescendants(subDescendants)
55+
expected := &Branch{
56+
Descendants: finalDescendants,
57+
}
58+
59+
assert.Equal(t, expected, branch)
60+
}

‎lib/trie/database.go

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,17 @@ func (t *Trie) load(db chaindb.Database, n Node) error {
219219
if err != nil {
220220
return fmt.Errorf("cannot load child at index %d with hash 0x%x: %w", i, hash, err)
221221
}
222+
223+
if decodedNode.Type() != node.LeafType { // branch decoded node
224+
// Note 1: the node is fully loaded with all its descendants
225+
// count only after the database load above.
226+
// Note 2: direct child node is already counted as descendant
227+
// when it was read as a leaf with hash only in decodeBranch,
228+
// so we only add the descendants of the child branch to the
229+
// current branch.
230+
childBranchDescendants := decodedNode.(*node.Branch).Descendants
231+
branch.AddDescendants(childBranchDescendants)
232+
}
222233
}
223234

224235
for _, key := range t.GetKeysWithPrefix(ChildStorageKeyPrefix) {

‎lib/trie/print_test.go

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,9 @@ func Test_Trie_String(t *testing.T) {
3939
"branch root": {
4040
trie: Trie{
4141
root: &node.Branch{
42-
Key: nil,
43-
Value: []byte{1, 2},
42+
Key: nil,
43+
Value: []byte{1, 2},
44+
Descendants: 2,
4445
Children: [16]node.Node{
4546
&node.Leaf{
4647
Key: []byte{1, 2, 3},
@@ -61,6 +62,7 @@ func Test_Trie_String(t *testing.T) {
6162
├── Dirty: false
6263
├── Key: nil
6364
├── Value: 0x0102
65+
├── Descendants: 2
6466
├── Calculated encoding: nil
6567
├── Calculated digest: nil
6668
├── Child 0

0 commit comments

Comments
 (0)