Skip to content

Commit 687045a

Browse files
fxamackerfxamacker
authored
Merge pull request #3339 from onflow/fxamacker/port-3050
Backport PR 3050 to v0.27 to reduce memory and speedup checkpoint by splitting trie into subtries
2 parents e6dbdae + 9f84ee5 commit 687045a

File tree

4 files changed

+607
-55
lines changed

4 files changed

+607
-55
lines changed

‎ledger/complete/mtrie/flattener/iterator.go

+5-6
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ package flattener
33
import (
44
"github.com/onflow/flow-go/ledger"
55
"github.com/onflow/flow-go/ledger/complete/mtrie/node"
6-
"github.com/onflow/flow-go/ledger/complete/mtrie/trie"
76
)
87

98
// NodeIterator is an iterator over the nodes in a trie.
@@ -76,13 +75,13 @@ type NodeIterator struct {
7675
// as for each node, the children have been previously encountered.
7776
// NodeIterator created by NewNodeIterator is safe for concurrent use
7877
// because visitedNodes is always nil in this case.
79-
func NewNodeIterator(mTrie *trie.MTrie) *NodeIterator {
78+
func NewNodeIterator(n *node.Node) *NodeIterator {
8079
// for a Trie with height H (measured by number of edges), the longest possible path contains H+1 vertices
8180
stackSize := ledger.NodeMaxHeight + 1
8281
i := &NodeIterator{
8382
stack: make([]*node.Node, 0, stackSize),
8483
}
85-
i.unprocessedRoot = mTrie.RootNode()
84+
i.unprocessedRoot = n
8685
return i
8786
}
8887

@@ -98,15 +97,15 @@ func NewNodeIterator(mTrie *trie.MTrie) *NodeIterator {
9897
// When re-building the Trie from the sequence of nodes, one can build the trie on the fly,
9998
// as for each node, the children have been previously encountered.
10099
// WARNING: visitedNodes is not safe for concurrent use.
101-
func NewUniqueNodeIterator(mTrie *trie.MTrie, visitedNodes map[*node.Node]uint64) *NodeIterator {
100+
func NewUniqueNodeIterator(n *node.Node, visitedNodes map[*node.Node]uint64) *NodeIterator {
102101
// For a Trie with height H (measured by number of edges), the longest possible path
103102
// contains H+1 vertices.
104103
stackSize := ledger.NodeMaxHeight + 1
105104
i := &NodeIterator{
106105
stack: make([]*node.Node, 0, stackSize),
107106
visitedNodes: visitedNodes,
108107
}
109-
i.unprocessedRoot = mTrie.RootNode()
108+
i.unprocessedRoot = n
110109
return i
111110
}
112111

@@ -115,7 +114,7 @@ func (i *NodeIterator) Next() bool {
115114
// initial call to Next() for a non-empty trie
116115
i.dig(i.unprocessedRoot)
117116
i.unprocessedRoot = nil
118-
return true
117+
return len(i.stack) > 0
119118
}
120119

121120
// the current head of the stack, `n`, has been recalled

‎ledger/complete/mtrie/flattener/iterator_test.go

+140-11
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ import (
1616
func TestEmptyTrie(t *testing.T) {
1717
emptyTrie := trie.NewEmptyMTrie()
1818

19-
itr := flattener.NewNodeIterator(emptyTrie)
19+
itr := flattener.NewNodeIterator(emptyTrie.RootNode())
2020
require.True(t, nil == itr.Value()) // initial iterator should return nil
2121

2222
require.False(t, itr.Next())
@@ -30,7 +30,7 @@ func TestPopulatedTrie(t *testing.T) {
3030
emptyTrie := trie.NewEmptyMTrie()
3131

3232
// key: 0000...
33-
p1 := testutils.PathByUint8(1)
33+
p1 := testutils.PathByUint8(0)
3434
v1 := testutils.LightPayload8('A', 'a')
3535

3636
// key: 0100....
@@ -43,12 +43,12 @@ func TestPopulatedTrie(t *testing.T) {
4343
testTrie, _, err := trie.NewTrieWithUpdatedRegisters(emptyTrie, paths, payloads, true)
4444
require.NoError(t, err)
4545

46-
for itr := flattener.NewNodeIterator(testTrie); itr.Next(); {
46+
for itr := flattener.NewNodeIterator(testTrie.RootNode()); itr.Next(); {
4747
fmt.Println(itr.Value().FmtStr("", ""))
4848
fmt.Println()
4949
}
5050

51-
itr := flattener.NewNodeIterator(testTrie)
51+
itr := flattener.NewNodeIterator(testTrie.RootNode())
5252

5353
require.True(t, itr.Next())
5454
p1_leaf := itr.Value()
@@ -80,13 +80,13 @@ func TestUniqueNodeIterator(t *testing.T) {
8080
emptyTrie := trie.NewEmptyMTrie()
8181

8282
// visitedNodes is nil
83-
itr := flattener.NewUniqueNodeIterator(emptyTrie, nil)
83+
itr := flattener.NewUniqueNodeIterator(emptyTrie.RootNode(), nil)
8484
require.False(t, itr.Next())
8585
require.True(t, nil == itr.Value()) // initial iterator should return nil
8686

8787
// visitedNodes is empty map
8888
visitedNodes := make(map[*node.Node]uint64)
89-
itr = flattener.NewUniqueNodeIterator(emptyTrie, visitedNodes)
89+
itr = flattener.NewUniqueNodeIterator(emptyTrie.RootNode(), visitedNodes)
9090
require.False(t, itr.Next())
9191
require.True(t, nil == itr.Value()) // initial iterator should return nil
9292
})
@@ -95,7 +95,7 @@ func TestUniqueNodeIterator(t *testing.T) {
9595
emptyTrie := trie.NewEmptyMTrie()
9696

9797
// key: 0000...
98-
p1 := testutils.PathByUint8(1)
98+
p1 := testutils.PathByUint8(0)
9999
v1 := testutils.LightPayload8('A', 'a')
100100

101101
// key: 0100....
@@ -126,7 +126,7 @@ func TestUniqueNodeIterator(t *testing.T) {
126126

127127
// visitedNodes is nil
128128
i := 0
129-
for itr := flattener.NewUniqueNodeIterator(updatedTrie, nil); itr.Next(); {
129+
for itr := flattener.NewUniqueNodeIterator(updatedTrie.RootNode(), nil); itr.Next(); {
130130
n := itr.Value()
131131
require.True(t, i < len(expectedNodes))
132132
require.Equal(t, expectedNodes[i], n)
@@ -138,7 +138,7 @@ func TestUniqueNodeIterator(t *testing.T) {
138138
// there isn't any shared sub-trie.
139139
visitedNodes := make(map[*node.Node]uint64)
140140
i = 0
141-
for itr := flattener.NewUniqueNodeIterator(updatedTrie, visitedNodes); itr.Next(); {
141+
for itr := flattener.NewUniqueNodeIterator(updatedTrie.RootNode(), visitedNodes); itr.Next(); {
142142
n := itr.Value()
143143
visitedNodes[n] = uint64(i)
144144

@@ -157,7 +157,7 @@ func TestUniqueNodeIterator(t *testing.T) {
157157
emptyTrie := trie.NewEmptyMTrie()
158158

159159
// key: 0000...
160-
p1 := testutils.PathByUint8(1)
160+
p1 := testutils.PathByUint8(0)
161161
v1 := testutils.LightPayload8('A', 'a')
162162

163163
// key: 0100....
@@ -254,7 +254,7 @@ func TestUniqueNodeIterator(t *testing.T) {
254254
visitedNodes := make(map[*node.Node]uint64)
255255
i := 0
256256
for _, trie := range tries {
257-
for itr := flattener.NewUniqueNodeIterator(trie, visitedNodes); itr.Next(); {
257+
for itr := flattener.NewUniqueNodeIterator(trie.RootNode(), visitedNodes); itr.Next(); {
258258
n := itr.Value()
259259
visitedNodes[n] = uint64(i)
260260

@@ -265,4 +265,133 @@ func TestUniqueNodeIterator(t *testing.T) {
265265
}
266266
require.Equal(t, i, len(expectedNodes))
267267
})
268+
269+
t.Run("subtries", func(t *testing.T) {
270+
271+
emptyTrie := trie.NewEmptyMTrie()
272+
273+
// key: 0000...
274+
p1 := testutils.PathByUint8(0)
275+
v1 := testutils.LightPayload8('A', 'a')
276+
277+
// key: 0100....
278+
p2 := testutils.PathByUint8(64)
279+
v2 := testutils.LightPayload8('B', 'b')
280+
281+
paths := []ledger.Path{p1, p2}
282+
payloads := []ledger.Payload{*v1, *v2}
283+
284+
trie1, _, err := trie.NewTrieWithUpdatedRegisters(emptyTrie, paths, payloads, true)
285+
require.NoError(t, err)
286+
287+
// trie1
288+
// n4
289+
// /
290+
// /
291+
// n3
292+
// / \
293+
// / \
294+
// n1 (p1/v1) n2 (p2/v2)
295+
//
296+
297+
// New trie reuses its parent's left sub-trie.
298+
299+
// key: 1000...
300+
p3 := testutils.PathByUint8(128)
301+
v3 := testutils.LightPayload8('C', 'c')
302+
303+
// key: 1100....
304+
p4 := testutils.PathByUint8(192)
305+
v4 := testutils.LightPayload8('D', 'd')
306+
307+
paths = []ledger.Path{p3, p4}
308+
payloads = []ledger.Payload{*v3, *v4}
309+
310+
trie2, _, err := trie.NewTrieWithUpdatedRegisters(trie1, paths, payloads, true)
311+
require.NoError(t, err)
312+
313+
// trie2
314+
// n8
315+
// / \
316+
// / \
317+
// n3 n7
318+
// (shared) / \
319+
// / \
320+
// n5 n6
321+
// (p3/v3) (p4/v4)
322+
323+
// New trie reuses its parent's right sub-trie, and left sub-trie's leaf node.
324+
325+
// key: 0000...
326+
v5 := testutils.LightPayload8('E', 'e')
327+
328+
paths = []ledger.Path{p1}
329+
payloads = []ledger.Payload{*v5}
330+
331+
trie3, _, err := trie.NewTrieWithUpdatedRegisters(trie2, paths, payloads, true)
332+
require.NoError(t, err)
333+
334+
// trie3
335+
// n11
336+
// / \
337+
// / \
338+
// n10 n7
339+
// / \ (shared)
340+
// / \
341+
// n9 n2
342+
// (p1/v5) (shared)
343+
344+
leftSubtries := []*node.Node{
345+
trie1.RootNode().LeftChild(),
346+
trie2.RootNode().LeftChild(),
347+
trie3.RootNode().LeftChild(),
348+
}
349+
350+
expectedNodesInLeftSubtries := []*node.Node{
351+
// unique nodes from trie1
352+
trie1.RootNode().LeftChild().LeftChild(), // n1
353+
trie1.RootNode().LeftChild().RightChild(), // n2
354+
trie1.RootNode().LeftChild(), // n3
355+
// unique nodes from trie3
356+
trie3.RootNode().LeftChild().LeftChild(), // n9
357+
trie3.RootNode().LeftChild(), // n10
358+
}
359+
360+
rightSubtries := []*node.Node{
361+
trie1.RootNode().RightChild(),
362+
trie2.RootNode().RightChild(),
363+
trie3.RootNode().RightChild(),
364+
}
365+
366+
expectedNodesInRightSubtries := []*node.Node{
367+
// unique nodes from trie2
368+
trie2.RootNode().RightChild().LeftChild(), // n5
369+
trie2.RootNode().RightChild().RightChild(), // n6
370+
trie2.RootNode().RightChild(), // n7
371+
}
372+
373+
testcases := []struct {
374+
roots []*node.Node
375+
expectedNodes []*node.Node
376+
}{
377+
{leftSubtries, expectedNodesInLeftSubtries},
378+
{rightSubtries, expectedNodesInRightSubtries},
379+
}
380+
381+
for _, tc := range testcases {
382+
visitedNodes := make(map[*node.Node]uint64)
383+
i := 0
384+
for _, n := range tc.roots {
385+
for itr := flattener.NewUniqueNodeIterator(n, visitedNodes); itr.Next(); {
386+
n := itr.Value()
387+
visitedNodes[n] = uint64(i)
388+
389+
require.True(t, i < len(tc.expectedNodes))
390+
require.Equal(t, tc.expectedNodes[i], n)
391+
i++
392+
}
393+
}
394+
require.Equal(t, i, len(tc.expectedNodes))
395+
}
396+
})
268397
}

0 commit comments

Comments
 (0)