-
Notifications
You must be signed in to change notification settings - Fork 175
/
trie.go
390 lines (353 loc) · 13.7 KB
/
trie.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
package trie
import (
"bytes"
"encoding/hex"
"errors"
"fmt"
"sync"
"github.com/hashicorp/go-multierror"
"github.com/onflow/flow-go/ledger"
"github.com/onflow/flow-go/ledger/common"
"github.com/onflow/flow-go/ledger/common/utils"
"github.com/onflow/flow-go/ledger/complete/mtrie/node"
)
// MTrie represents a perfect in-memory full binary Merkle tree with uniform height.
// For a detailed description of the storage model, please consult `mtrie/README.md`
//
// A MTrie is a thin wrapper around a the trie's root Node. An MTrie implements the
// logic for forming MTrie-graphs from the elementary nodes. Specifically:
// * how Nodes (graph vertices) form a Trie,
// * how register values are read from the trie,
// * how Merkle proofs are generated from a trie, and
// * how a new Trie with updated values is generated.
//
// `MTrie`s are _immutable_ data structures. Updating register values is implemented through
// copy-on-write, which creates a new `MTrie`. For minimal memory consumption, all sub-tries
// that where not affected by the write operation are shared between the original MTrie
// (before the register updates) and the updated MTrie (after the register writes).
//
// DEFINITIONS and CONVENTIONS:
// * HEIGHT of a node v in a tree is the number of edges on the longest downward path
// between v and a tree leaf. The height of a tree is the heights of its root.
// The height of a Trie is always the height of the fully-expanded tree.
type MTrie struct {
root *node.Node
height int
pathByteSize int
}
// NewEmptyMTrie returns an empty Mtrie (root is an empty node)
func NewEmptyMTrie(pathByteSize int) (*MTrie, error) {
if pathByteSize < 1 {
return nil, errors.New("trie's path size [in bytes] must be positive")
}
height := pathByteSize * 8
return &MTrie{
root: node.NewEmptyTreeRoot(height),
pathByteSize: pathByteSize,
height: height,
}, nil
}
// NewMTrie returns a Mtrie given the root
func NewMTrie(root *node.Node) (*MTrie, error) {
if root.Height()%8 != 0 {
return nil, errors.New("height of root node must be integer-multiple of 8")
}
pathByteSize := root.Height() / 8
return &MTrie{
root: root,
pathByteSize: pathByteSize,
height: root.Height(),
}, nil
}
// StringRootHash returns the trie's Hex-encoded root hash.
// Concurrency safe (as Tries are immutable structures by convention)
func (mt *MTrie) StringRootHash() string { return hex.EncodeToString(mt.root.Hash()) }
// RootHash returns the trie's root hash (i.e. the hash of the trie's root node).
// Concurrency safe (as Tries are immutable structures by convention)
func (mt *MTrie) RootHash() []byte { return mt.root.Hash() }
// PathLength return the length [in bytes] the trie operates with.
// Concurrency safe (as Tries are immutable structures by convention)
func (mt *MTrie) PathLength() int { return mt.pathByteSize }
// AllocatedRegCount returns the number of allocated registers in the trie.
// Concurrency safe (as Tries are immutable structures by convention)
func (mt *MTrie) AllocatedRegCount() uint64 { return mt.root.RegCount() }
// MaxDepth returns the length of the longest branch from root to leaf.
// Concurrency safe (as Tries are immutable structures by convention)
func (mt *MTrie) MaxDepth() uint16 { return mt.root.MaxDepth() }
// RootNode returns the Trie's root Node
// Concurrency safe (as Tries are immutable structures by convention)
func (mt *MTrie) RootNode() *node.Node {
return mt.root
}
// StringRootHash returns the trie's string representation.
// Concurrency safe (as Tries are immutable structures by convention)
func (mt *MTrie) String() string {
trieStr := fmt.Sprintf("Trie root hash: %v\n", mt.StringRootHash())
return trieStr + mt.root.FmtStr("", "")
}
// UnsafeRead read payloads for the given paths. It is called unsafe as it requires the
// paths to be sorted
// TODO move consistency checks from Forrest into Trie to obtain a safe, self-contained API
func (mt *MTrie) UnsafeRead(paths []ledger.Path) ([]*ledger.Payload, error) {
return mt.read(mt.root, paths)
}
func (mt *MTrie) read(head *node.Node, paths []ledger.Path) ([]*ledger.Payload, error) {
// path not found
if head == nil {
res := make([]*ledger.Payload, 0, len(paths))
for range paths {
res = append(res, ledger.EmptyPayload())
}
return res, nil
}
// reached a leaf node
if head.IsLeaf() {
res := make([]*ledger.Payload, 0)
for _, p := range paths {
if bytes.Equal(head.Path(), p) {
res = append(res, head.Payload())
} else {
res = append(res, ledger.EmptyPayload())
}
}
return res, nil
}
lpaths, rpaths, err := utils.SplitSortedPaths(paths, mt.height-head.Height())
if err != nil {
return nil, fmt.Errorf("can't read due to split path error: %w", err)
}
// TODO make this parallel
payloads := make([]*ledger.Payload, 0)
if len(lpaths) > 0 {
p, err := mt.read(head.LeftChild(), lpaths)
if err != nil {
return nil, err
}
payloads = append(payloads, p...)
}
if len(rpaths) > 0 {
p, err := mt.read(head.RigthChild(), rpaths)
if err != nil {
return nil, err
}
payloads = append(payloads, p...)
}
return payloads, nil
}
// NewTrieWithUpdatedRegisters constructs a new trie containing all registers from the parent trie.
// The key-value pairs specify the registers whose values are supposed to hold updated values
// compared to the parent trie. Constructing the new trie is done in a COPY-ON-WRITE manner:
// * The original trie remains unchanged.
// * subtries that remain unchanged are from the parent trie instead of copied.
// UNSAFE: method requires the following conditions to be satisfied:
// * keys are NOT duplicated
// TODO: move consistency checks from MForest to here, to make API is safe and self-contained
func NewTrieWithUpdatedRegisters(parentTrie *MTrie, updatedPaths []ledger.Path, updatedPayloads []ledger.Payload) (*MTrie, error) {
parentRoot := parentTrie.root
updatedRoot, err := update(parentTrie.height, parentRoot.Height(), parentRoot, updatedPaths, updatedPayloads)
if err != nil {
return nil, fmt.Errorf("constructing updated trie failed: %w", err)
}
updatedTrie, err := NewMTrie(updatedRoot)
if err != nil {
return nil, fmt.Errorf("constructing updated trie failed: %w", err)
}
return updatedTrie, nil
}
// update returns the head of updated sub-trie for the specified key-value pairs.
// UNSAFE: update requires the following conditions to be satisfied,
// but does not explicitly check them for performance reasons
// * all keys AND the parent node share the same common prefix [0 : mt.maxHeight-1 - headHeight)
// (excluding the bit at index headHeight)
// * keys are NOT duplicated
// TODO: remove error return
func update(treeHeight int, nodeHeight int, parentNode *node.Node, paths []ledger.Path, payloads []ledger.Payload) (*node.Node, error) {
if parentNode == nil { // parent Trie has no sub-trie for the set of paths => construct entire subtree
return constructSubtrie(treeHeight, nodeHeight, paths, payloads)
}
if len(paths) == 0 { // We are not changing any values in this sub-trie => return parent trie
return parentNode, nil
}
// from here on, we have parentNode != nil AND len(paths) > 0
if parentNode.IsLeaf() { // parent node is a leaf, i.e. parent Trie only stores a single value in this sub-trie
parentPath := parentNode.Path() // Per definition, a leaf must have a non-nil path
overrideExistingValue := false // true if and only if we are updating the parent Trie's leaf node value
for _, p := range paths {
if bytes.Equal(p, parentPath) {
overrideExistingValue = true
break
}
}
if !overrideExistingValue {
// TODO: copy payload when using in-place MergeSort for separating the payloads
paths = append(paths, parentNode.Path())
payloads = append(payloads, *parentNode.Payload())
}
return constructSubtrie(treeHeight, nodeHeight, paths, payloads)
}
// Split payloads so we can update the trie in parallel
lpaths, lpayloads, rpaths, rpayloads, err := utils.SplitByPath(paths, payloads, treeHeight-nodeHeight)
if err != nil {
return nil, fmt.Errorf("error spliting payloads by path: %w", err)
}
// TODO [runtime optimization]: do not branch if either lpayload or rpayload is empty
var lChild, rChild *node.Node
var lErr, rErr error
wg := sync.WaitGroup{}
wg.Add(1)
go func() {
defer wg.Done()
lChild, lErr = update(treeHeight, nodeHeight-1, parentNode.LeftChild(), lpaths, lpayloads)
}()
rChild, rErr = update(treeHeight, nodeHeight-1, parentNode.RigthChild(), rpaths, rpayloads)
wg.Wait()
if lErr != nil || rErr != nil {
var merr *multierror.Error
if lErr != nil {
merr = multierror.Append(merr, lErr)
}
if rErr != nil {
merr = multierror.Append(merr, rErr)
}
if err := merr.ErrorOrNil(); err != nil {
return nil, fmt.Errorf("internal error while updating trie: %w", err)
}
}
return node.NewInterimNode(nodeHeight, lChild, rChild), nil
}
// constructSubtrie returns the head of a newly-constructed sub-trie for the specified key-value pairs.
// UNSAFE: constructSubtrie requires the following conditions to be satisfied,
// but does not explicitly check them for performance reasons
// * paths all share the same common prefix [0 : mt.maxHeight-1 - headHeight)
// (excluding the bit at index headHeight)
// * paths contains at least one element
// * paths are NOT duplicated
// TODO: remove error return
func constructSubtrie(treeHeight int, nodeHeight int, paths []ledger.Path, payloads []ledger.Payload) (*node.Node, error) {
// no inserts => default value, represented by nil node
if len(paths) == 0 {
return nil, nil
}
// If we are at a leaf node, we create the node
if len(paths) == 1 {
return node.NewLeaf(paths[0], &payloads[0], nodeHeight), nil
}
// from here on, we have: len(paths) > 1
// Split updates by paths so we can update the trie in parallel
lpaths, lpayloads, rpaths, rpayloads, err := utils.SplitByPath(paths, payloads, treeHeight-nodeHeight)
// Note: (pathLength-height) will never reach the value pathLength, i.e. we will never execute this code for height==0
// This is because at height=0, we only have (at most) one path left, as paths are not duplicated
// (by requirement of this function). But even if this condition is violated, the code will not return a faulty
// but instead panic with Index Out Of Range error
if err != nil {
return nil, fmt.Errorf("error spliting paths: %w", err)
}
// TODO [runtime optimization]: do not branch if either lpaths or rpaths is empty
var lChild, rChild *node.Node
var lErr, rErr error
wg := sync.WaitGroup{}
wg.Add(1)
go func() {
defer wg.Done()
lChild, lErr = constructSubtrie(treeHeight, nodeHeight-1, lpaths, lpayloads)
}()
rChild, rErr = constructSubtrie(treeHeight, nodeHeight-1, rpaths, rpayloads)
wg.Wait()
if lErr != nil || rErr != nil {
var merr *multierror.Error
if lErr != nil {
merr = multierror.Append(merr, lErr)
}
if rErr != nil {
merr = multierror.Append(merr, rErr)
}
if err := merr.ErrorOrNil(); err != nil {
return nil, fmt.Errorf("internal error while constructing sub-trie: %w", err)
}
}
return node.NewInterimNode(nodeHeight, lChild, rChild), nil
}
// UnsafeProofs provides proofs for the given paths, this is called unsafe as
// it requires the input paths to be sorted in advance.
func (mt *MTrie) UnsafeProofs(paths []ledger.Path, proofs []*ledger.TrieProof) error {
return mt.proofs(mt.root, paths, proofs)
}
func (mt *MTrie) proofs(head *node.Node, paths []ledger.Path, proofs []*ledger.TrieProof) error {
// we've reached the end of a trie
// and path is not found (noninclusion proof)
if head == nil {
return nil
}
// we've reached a leaf
if head.IsLeaf() {
// value matches (inclusion proof)
if bytes.Equal(head.Path(), paths[0]) {
proofs[0].Path = head.Path()
proofs[0].Payload = head.Payload()
proofs[0].Inclusion = true
}
// TODO: insert ERROR if len(paths) != 1
return nil
}
// increment steps for all the proofs
for _, p := range proofs {
p.Steps++
}
// split paths based on the value of i-th bit (i = trie height - node height)
lpaths, lproofs, rpaths, rproofs, err := utils.SplitTrieProofsByPath(paths, proofs, mt.height-head.Height())
if err != nil {
return fmt.Errorf("proof generation failed, path split error: %w", err)
}
if len(lpaths) > 0 {
if rChild := head.RigthChild(); rChild != nil {
nodeHash := rChild.Hash()
isDef := bytes.Equal(nodeHash, common.GetDefaultHashForHeight(rChild.Height()))
if !isDef { // in proofs, we only provide non-default value hashes
for _, p := range lproofs {
err := utils.SetBit(p.Flags, mt.height-head.Height())
if err != nil {
return err
}
p.Interims = append(p.Interims, nodeHash)
}
}
}
err := mt.proofs(head.LeftChild(), lpaths, lproofs)
if err != nil {
return err
}
}
if len(rpaths) > 0 {
if lChild := head.LeftChild(); lChild != nil {
nodeHash := lChild.Hash()
isDef := bytes.Equal(nodeHash, common.GetDefaultHashForHeight(lChild.Height()))
if !isDef { // in proofs, we only provide non-default value hashes
for _, p := range rproofs {
err := utils.SetBit(p.Flags, mt.height-head.Height())
if err != nil {
return err
}
p.Interims = append(p.Interims, nodeHash)
}
}
}
err := mt.proofs(head.RigthChild(), rpaths, rproofs)
if err != nil {
return err
}
}
return nil
}
// Equals compares two tries for equality.
// Tries are equal iff they store the same data (i.e. root hash matches)
// and their number and height are identical
func (mt *MTrie) Equals(o *MTrie) bool {
if o == nil {
return false
}
return o.PathLength() == mt.PathLength() && bytes.Equal(o.RootHash(), mt.RootHash())
}
// EmptyTrieRootHash returns the rootHash of an empty Trie for the specified path size [bytes]
func EmptyTrieRootHash(pathByteSize int) []byte {
return node.NewEmptyTreeRoot(8 * pathByteSize).Hash()
}