arbo/tree.go
tree.goRaw
1/*
2Package arbo implements a Merkle Tree compatible with the circomlib
3implementation of the MerkleTree, following the specification from
4https://docs.iden3.io/publications/pdfs/Merkle-Tree.pdf and
5https://eprint.iacr.org/2018/955.
6
7Allows to define which hash function to use. So for example, when working with
8zkSnarks the Poseidon hash function can be used, but when not, it can be used
9the Blake2b hash function, which has much faster computation time.
10*/
11package arbo
12
13import (
14 "bytes"
15 "encoding/binary"
16 "encoding/hex"
17 "fmt"
18 "io"
19 "math"
20 "runtime"
21 "sync"
22
23 "go.vocdoni.io/dvote/db"
24)
25
26const (
27 // PrefixValueLen defines the bytes-prefix length used for the Value
28 // bytes representation stored in the db
29 PrefixValueLen = 2
30
31 // PrefixValueEmpty is used for the first byte of a Value to indicate
32 // that is an Empty value
33 PrefixValueEmpty = 0
34 // PrefixValueLeaf is used for the first byte of a Value to indicate
35 // that is a Leaf value
36 PrefixValueLeaf = 1
37 // PrefixValueIntermediate is used for the first byte of a Value to
38 // indicate that is a Intermediate value
39 PrefixValueIntermediate = 2
40
41 // nChars is used to crop the Graphviz nodes labels
42 nChars = 4
43
44 maxUint8 = int(^uint8(0)) // 2**8 -1
45 maxUint16 = int(^uint16(0)) // 2**16 -1
46)
47
48var (
49 // DefaultThresholdNLeafs defines the threshold number of leafs in the
50 // tree that determines if AddBatch will work in memory or in disk. It
51 // is defined when calling NewTree, and if set to 0 it will work always
52 // in disk.
53 DefaultThresholdNLeafs = 65536
54
55 dbKeyRoot = []byte("root")
56 dbKeyNLeafs = []byte("nleafs")
57 emptyValue = []byte{0}
58
59 // ErrKeyNotFound is used when a key is not found in the db neither in
60 // the current db Batch.
61 ErrKeyNotFound = fmt.Errorf("key not found")
62 // ErrKeyAlreadyExists is used when trying to add a key as leaf to the
63 // tree that already exists.
64 ErrKeyAlreadyExists = fmt.Errorf("key already exists")
65 // ErrInvalidValuePrefix is used when going down into the tree, a value
66 // is read from the db and has an unrecognized prefix.
67 ErrInvalidValuePrefix = fmt.Errorf("invalid value prefix")
68 // ErrDBNoTx is used when trying to use Tree.dbPut but Tree.dbBatch==nil
69 ErrDBNoTx = fmt.Errorf("dbPut error: no db Batch")
70 // ErrMaxLevel indicates when going down into the tree, the max level is
71 // reached
72 ErrMaxLevel = fmt.Errorf("max level reached")
73 // ErrMaxVirtualLevel indicates when going down into the tree, the max
74 // virtual level is reached
75 ErrMaxVirtualLevel = fmt.Errorf("max virtual level reached")
76 // ErrSnapshotNotEditable indicates when the tree is a non writable
77 // snapshot, thus can not be modified
78 ErrSnapshotNotEditable = fmt.Errorf("snapshot tree can not be edited")
79 // ErrTreeNotEmpty indicates when the tree was expected to be empty and
80 // it is not
81 ErrTreeNotEmpty = fmt.Errorf("tree is not empty")
82)
83
84// Tree defines the struct that implements the MerkleTree functionalities
85type Tree struct {
86 sync.Mutex
87
88 db db.Database
89 maxLevels int
90 // thresholdNLeafs defines the threshold number of leafs in the tree
91 // that determines if AddBatch will work in memory or in disk. It is
92 // defined when calling NewTree, and if set to 0 it will work always in
93 // disk.
94 thresholdNLeafs int
95 snapshotRoot []byte
96
97 hashFunction HashFunction
98 // TODO in the methods that use it, check if emptyHash param is len>0
99 // (check if it has been initialized)
100 emptyHash []byte
101
102 dbg *dbgStats
103}
104
105// Config defines the configuration for calling NewTree & NewTreeWithTx methods
106type Config struct {
107 Database db.Database
108 MaxLevels int
109 ThresholdNLeafs int
110 HashFunction HashFunction
111}
112
113// NewTree returns a new Tree, if there is a Tree still in the given database, it
114// will load it.
115func NewTree(cfg Config) (*Tree, error) {
116 wTx := cfg.Database.WriteTx()
117 defer wTx.Discard()
118
119 t, err := NewTreeWithTx(wTx, cfg)
120 if err != nil {
121 return nil, err
122 }
123
124 if err = wTx.Commit(); err != nil {
125 return nil, err
126 }
127 return t, nil
128}
129
130// NewTreeWithTx returns a new Tree using the given db.WriteTx, which will not
131// be ccommited inside this method, if there is a Tree still in the given
132// database, it will load it.
133func NewTreeWithTx(wTx db.WriteTx, cfg Config) (*Tree, error) {
134 // if thresholdNLeafs is set to 0, use the DefaultThresholdNLeafs
135 t := Tree{db: cfg.Database, maxLevels: cfg.MaxLevels,
136 thresholdNLeafs: cfg.ThresholdNLeafs, hashFunction: cfg.HashFunction}
137 t.emptyHash = make([]byte, t.hashFunction.Len()) // empty
138
139 _, err := wTx.Get(dbKeyRoot)
140 if err == db.ErrKeyNotFound {
141 // store new root 0 (empty)
142 if err = wTx.Set(dbKeyRoot, t.emptyHash); err != nil {
143 return nil, err
144 }
145 if err = t.setNLeafs(wTx, 0); err != nil {
146 return nil, err
147 }
148 return &t, nil
149 } else if err != nil {
150 return nil, err
151 }
152 return &t, nil
153}
154
155// Root returns the root of the Tree
156func (t *Tree) Root() ([]byte, error) {
157 rTx := t.db.ReadTx()
158 defer rTx.Discard()
159 return t.RootWithTx(rTx)
160}
161
162// RootWithTx returns the root of the Tree using the given db.ReadTx
163func (t *Tree) RootWithTx(rTx db.ReadTx) ([]byte, error) {
164 // if snapshotRoot is defined, means that the tree is a snapshot, and
165 // the root is not obtained from the db, but from the snapshotRoot
166 // parameter
167 if t.snapshotRoot != nil {
168 return t.snapshotRoot, nil
169 }
170 // get db root
171 return rTx.Get(dbKeyRoot)
172}
173
174func (t *Tree) setRoot(wTx db.WriteTx, root []byte) error {
175 return wTx.Set(dbKeyRoot, root)
176}
177
178// HashFunction returns Tree.hashFunction
179func (t *Tree) HashFunction() HashFunction {
180 return t.hashFunction
181}
182
183// editable returns true if the tree is editable, and false when is not
184// editable (because is a snapshot tree)
185func (t *Tree) editable() bool {
186 return t.snapshotRoot == nil
187}
188
189// Invalid is used when a key-value can not be added trough AddBatch, and
190// contains the index of the key-value and the error.
191type Invalid struct {
192 Index int
193 Error error
194}
195
196// AddBatch adds a batch of key-values to the Tree. Returns an array containing
197// the indexes of the keys failed to add. Supports empty values as input
198// parameters, which is equivalent to 0 valued byte array.
199func (t *Tree) AddBatch(keys, values [][]byte) ([]Invalid, error) {
200 wTx := t.db.WriteTx()
201 defer wTx.Discard()
202
203 invalids, err := t.AddBatchWithTx(wTx, keys, values)
204 if err != nil {
205 return invalids, err
206 }
207 return invalids, wTx.Commit()
208}
209
210// AddBatchWithTx does the same than the AddBatch method, but allowing to pass
211// the db.WriteTx that is used. The db.WriteTx will not be committed inside
212// this method.
213func (t *Tree) AddBatchWithTx(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
214 t.Lock()
215 defer t.Unlock()
216
217 if !t.editable() {
218 return nil, ErrSnapshotNotEditable
219 }
220
221 e := []byte{}
222 // equal the number of keys & values
223 if len(keys) > len(values) {
224 // add missing values
225 for i := len(values); i < len(keys); i++ {
226 values = append(values, e)
227 }
228 } else if len(keys) < len(values) {
229 // crop extra values
230 values = values[:len(keys)]
231 }
232
233 nLeafs, err := t.GetNLeafsWithTx(wTx)
234 if err != nil {
235 return nil, err
236 }
237 if nLeafs > t.thresholdNLeafs {
238 return t.addBatchInDisk(wTx, keys, values)
239 }
240 return t.addBatchInMemory(wTx, keys, values)
241}
242
243func (t *Tree) addBatchInDisk(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
244 nCPU := flp2(runtime.NumCPU())
245 if nCPU == 1 || len(keys) < nCPU {
246 var invalids []Invalid
247 for i := 0; i < len(keys); i++ {
248 if err := t.addWithTx(wTx, keys[i], values[i]); err != nil {
249 invalids = append(invalids, Invalid{i, err})
250 }
251 }
252 return invalids, nil
253 }
254
255 kvs, invalids, err := keysValuesToKvs(t.maxLevels, keys, values)
256 if err != nil {
257 return nil, err
258 }
259
260 buckets := splitInBuckets(kvs, nCPU)
261
262 root, err := t.RootWithTx(wTx)
263 if err != nil {
264 return nil, err
265 }
266
267 l := int(math.Log2(float64(nCPU)))
268 subRoots, err := t.getSubRootsAtLevel(wTx, root, l+1)
269 if err != nil {
270 return nil, err
271 }
272 if len(subRoots) != nCPU {
273 // Already populated Tree but Unbalanced.
274
275 // add one key at each bucket, and then continue with the flow
276 for i := 0; i < len(buckets); i++ {
277 // add one leaf of the bucket, if there is an error when
278 // adding the k-v, try to add the next one of the bucket
279 // (until one is added)
280 inserted := -1
281 for j := 0; j < len(buckets[i]); j++ {
282 if newRoot, err := t.add(wTx, root, 0,
283 buckets[i][j].k, buckets[i][j].v); err == nil {
284 inserted = j
285 root = newRoot
286 break
287 }
288 }
289
290 // remove the inserted element from buckets[i]
291 if inserted != -1 {
292 buckets[i] = append(buckets[i][:inserted], buckets[i][inserted+1:]...)
293 }
294 }
295 subRoots, err = t.getSubRootsAtLevel(wTx, root, l+1)
296 if err != nil {
297 return nil, err
298 }
299 }
300
301 if len(subRoots) != nCPU {
302 return nil, fmt.Errorf("This error should not be reached."+
303 " len(subRoots) != nCPU, len(subRoots)=%d, nCPU=%d."+
304 " Please report it in a new issue:"+
305 " https://github.com/vocdoni/arbo/issues/new", len(subRoots), nCPU)
306 }
307
308 invalidsInBucket := make([][]Invalid, nCPU)
309 txs := make([]db.WriteTx, nCPU)
310 for i := 0; i < nCPU; i++ {
311 txs[i] = t.db.WriteTx()
312 err := txs[i].Apply(wTx)
313 if err != nil {
314 return nil, err
315 }
316 }
317
318 var wg sync.WaitGroup
319 wg.Add(nCPU)
320 for i := 0; i < nCPU; i++ {
321 go func(cpu int) {
322 // use different wTx for each cpu, after once all
323 // are done, iter over the cpuWTxs and copy their
324 // content into the main wTx
325 for j := 0; j < len(buckets[cpu]); j++ {
326 newSubRoot, err := t.add(txs[cpu], subRoots[cpu],
327 l, buckets[cpu][j].k, buckets[cpu][j].v)
328 if err != nil {
329 invalidsInBucket[cpu] = append(invalidsInBucket[cpu],
330 Invalid{buckets[cpu][j].pos, err})
331 continue
332 }
333 // if there has not been errors, set the new subRoots[cpu]
334 subRoots[cpu] = newSubRoot
335 }
336 wg.Done()
337 }(i)
338 }
339 wg.Wait()
340
341 for i := 0; i < nCPU; i++ {
342 if err := wTx.Apply(txs[i]); err != nil {
343 return nil, err
344 }
345 txs[i].Discard()
346 }
347
348 for i := 0; i < len(invalidsInBucket); i++ {
349 invalids = append(invalids, invalidsInBucket[i]...)
350 }
351
352 newRoot, err := t.upFromSubRoots(wTx, subRoots)
353 if err != nil {
354 return nil, err
355 }
356
357 // update dbKeyNLeafs
358 if err := t.SetRootWithTx(wTx, newRoot); err != nil {
359 return nil, err
360 }
361
362 // update nLeafs
363 if err := t.incNLeafs(wTx, len(keys)-len(invalids)); err != nil {
364 return nil, err
365 }
366
367 return invalids, nil
368}
369
370func (t *Tree) upFromSubRoots(wTx db.WriteTx, subRoots [][]byte) ([]byte, error) {
371 // is a method of Tree just to get access to t.hashFunction and
372 // t.emptyHash.
373
374 // go up from subRoots to up, storing nodes in the given WriteTx
375 // once up at the root, store it in the WriteTx using the dbKeyRoot
376 if len(subRoots) == 1 {
377 return subRoots[0], nil
378 }
379 // get the subRoots values to know the node types of each subRoot
380 nodeTypes := make([]byte, len(subRoots))
381 for i := 0; i < len(subRoots); i++ {
382 if bytes.Equal(subRoots[i], t.emptyHash) {
383 nodeTypes[i] = PrefixValueEmpty
384 continue
385 }
386 v, err := wTx.Get(subRoots[i])
387 if err != nil {
388 return nil, err
389 }
390 nodeTypes[i] = v[0]
391 }
392
393 var newSubRoots [][]byte
394 for i := 0; i < len(subRoots); i += 2 {
395 if (bytes.Equal(subRoots[i], t.emptyHash) && bytes.Equal(subRoots[i+1], t.emptyHash)) ||
396 (nodeTypes[i] == PrefixValueLeaf && bytes.Equal(subRoots[i+1], t.emptyHash)) {
397 // when both sub nodes are empty, the parent is also empty
398 // or
399 // when 1st sub node is a leaf but the 2nd is empty, the
400 // leaf is used as 'parent'
401
402 newSubRoots = append(newSubRoots, subRoots[i])
403 continue
404 }
405 if bytes.Equal(subRoots[i], t.emptyHash) && nodeTypes[i+1] == PrefixValueLeaf {
406 // when 2nd sub node is a leaf but the 1st is empty,
407 // the leaf is used as 'parent'
408 newSubRoots = append(newSubRoots, subRoots[i+1])
409 continue
410 }
411
412 k, v, err := t.newIntermediate(subRoots[i], subRoots[i+1])
413 if err != nil {
414 return nil, err
415 }
416 // store k-v to db
417 if err = wTx.Set(k, v); err != nil {
418 return nil, err
419 }
420 newSubRoots = append(newSubRoots, k)
421 }
422
423 return t.upFromSubRoots(wTx, newSubRoots)
424}
425
426func (t *Tree) getSubRootsAtLevel(rTx db.ReadTx, root []byte, l int) ([][]byte, error) {
427 // go at level l and return each node key, where each node key is the
428 // subRoot of the subTree that starts there
429
430 var subRoots [][]byte
431 err := t.iterWithStop(rTx, root, 0, func(currLvl int, k, v []byte) bool {
432 if currLvl == l && !bytes.Equal(k, t.emptyHash) {
433 subRoots = append(subRoots, k)
434 }
435 if currLvl >= l {
436 return true // to stop the iter from going down
437 }
438 return false
439 })
440
441 return subRoots, err
442}
443
444func (t *Tree) addBatchInMemory(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
445 vt, err := t.loadVT(wTx)
446 if err != nil {
447 return nil, err
448 }
449
450 invalids, err := vt.addBatch(keys, values)
451 if err != nil {
452 return nil, err
453 }
454
455 // once the VirtualTree is build, compute the hashes
456 pairs, err := vt.computeHashes()
457 if err != nil {
458 // currently invalids in computeHashes are not counted,
459 // but should not be needed, as if there is an error there is
460 // nothing stored in the db and the error is returned
461 return nil, err
462 }
463
464 // store pairs in db
465 for i := 0; i < len(pairs); i++ {
466 if err := wTx.Set(pairs[i][0], pairs[i][1]); err != nil {
467 return nil, err
468 }
469 }
470
471 // store root (from the vt) to db
472 if vt.root != nil {
473 if err := wTx.Set(dbKeyRoot, vt.root.h); err != nil {
474 return nil, err
475 }
476 }
477
478 // update nLeafs
479 if err := t.incNLeafs(wTx, len(keys)-len(invalids)); err != nil {
480 return nil, err
481 }
482
483 return invalids, nil
484}
485
486// loadVT loads a new virtual tree (vt) from the current Tree, which contains
487// the same leafs.
488func (t *Tree) loadVT(rTx db.ReadTx) (vt, error) {
489 vt := newVT(t.maxLevels, t.hashFunction)
490 vt.params.dbg = t.dbg
491 var callbackErr error
492 err := t.IterateWithStopWithTx(rTx, nil, func(_ int, k, v []byte) bool {
493 if v[0] != PrefixValueLeaf {
494 return false
495 }
496 leafK, leafV := ReadLeafValue(v)
497 if err := vt.add(0, leafK, leafV); err != nil {
498 callbackErr = err
499 return true
500 }
501 return false
502 })
503 if callbackErr != nil {
504 return vt, callbackErr
505 }
506
507 return vt, err
508}
509
510// Add inserts the key-value into the Tree. If the inputs come from a
511// *big.Int, is expected that are represented by a Little-Endian byte array
512// (for circom compatibility).
513func (t *Tree) Add(k, v []byte) error {
514 wTx := t.db.WriteTx()
515 defer wTx.Discard()
516
517 if err := t.AddWithTx(wTx, k, v); err != nil {
518 return err
519 }
520
521 return wTx.Commit()
522}
523
524// AddWithTx does the same than the Add method, but allowing to pass the
525// db.WriteTx that is used. The db.WriteTx will not be committed inside this
526// method.
527func (t *Tree) AddWithTx(wTx db.WriteTx, k, v []byte) error {
528 t.Lock()
529 defer t.Unlock()
530
531 if !t.editable() {
532 return ErrSnapshotNotEditable
533 }
534 return t.addWithTx(wTx, k, v)
535}
536
537// warning: addWithTx does not use the Tree mutex, the mutex is responsibility
538// of the methods calling this method, and same with t.editable().
539func (t *Tree) addWithTx(wTx db.WriteTx, k, v []byte) error {
540 root, err := t.RootWithTx(wTx)
541 if err != nil {
542 return err
543 }
544
545 root, err = t.add(wTx, root, 0, k, v) // add from level 0
546 if err != nil {
547 return err
548 }
549 // store root to db
550 if err := t.setRoot(wTx, root); err != nil {
551 return err
552 }
553 // update nLeafs
554 if err = t.incNLeafs(wTx, 1); err != nil {
555 return err
556 }
557 return nil
558}
559
560// keyPathFromKey returns the keyPath and checks that the key is not bigger
561// than maximum key length for the tree maxLevels size.
562// This is because if the key bits length is bigger than the maxLevels of the
563// tree, two different keys that their difference is at the end, will collision
564// in the same leaf of the tree (at the max depth).
565func keyPathFromKey(maxLevels int, k []byte) ([]byte, error) {
566 maxKeyLen := int(math.Ceil(float64(maxLevels) / float64(8))) //nolint:gomnd
567 if len(k) > maxKeyLen {
568 return nil, fmt.Errorf("len(k) can not be bigger than ceil(maxLevels/8), where"+
569 " len(k): %d, maxLevels: %d, max key len=ceil(maxLevels/8): %d. Might need"+
570 " a bigger tree depth (maxLevels>=%d) in order to input keys of length %d",
571 len(k), maxLevels, maxKeyLen, len(k)*8, len(k)) //nolint:gomnd
572 }
573 keyPath := make([]byte, maxKeyLen) //nolint:gomnd
574 copy(keyPath[:], k)
575 return keyPath, nil
576}
577
578func (t *Tree) add(wTx db.WriteTx, root []byte, fromLvl int, k, v []byte) ([]byte, error) {
579 keyPath, err := keyPathFromKey(t.maxLevels, k)
580 if err != nil {
581 return nil, err
582 }
583 path := getPath(t.maxLevels, keyPath)
584
585 // go down to the leaf
586 var siblings [][]byte
587 _, _, siblings, err = t.down(wTx, k, root, siblings, path, fromLvl, false)
588 if err != nil {
589 return nil, err
590 }
591
592 leafKey, leafValue, err := t.newLeafValue(k, v)
593 if err != nil {
594 return nil, err
595 }
596
597 if err := wTx.Set(leafKey, leafValue); err != nil {
598 return nil, err
599 }
600
601 // go up to the root
602 if len(siblings) == 0 {
603 // return the leafKey as root
604 return leafKey, nil
605 }
606 root, err = t.up(wTx, leafKey, siblings, path, len(siblings)-1, fromLvl)
607 if err != nil {
608 return nil, err
609 }
610
611 return root, nil
612}
613
614// down goes down to the leaf recursively
615func (t *Tree) down(rTx db.ReadTx, newKey, currKey []byte, siblings [][]byte,
616 path []bool, currLvl int, getLeaf bool) (
617 []byte, []byte, [][]byte, error) {
618 if currLvl > t.maxLevels {
619 return nil, nil, nil, ErrMaxLevel
620 }
621
622 var err error
623 var currValue []byte
624 if bytes.Equal(currKey, t.emptyHash) {
625 // empty value
626 return currKey, emptyValue, siblings, nil
627 }
628 currValue, err = rTx.Get(currKey)
629 if err != nil {
630 return nil, nil, nil, err
631 }
632
633 switch currValue[0] {
634 case PrefixValueEmpty: // empty
635 fmt.Printf("newKey: %s, currKey: %s, currLvl: %d, currValue: %s\n",
636 hex.EncodeToString(newKey), hex.EncodeToString(currKey),
637 currLvl, hex.EncodeToString(currValue))
638 panic("This point should not be reached, as the 'if currKey==t.emptyHash'" +
639 " above should avoid reaching this point. This panic is temporary" +
640 " for reporting purposes, will be deleted in future versions." +
641 " Please paste this log (including the previous log lines) in a" +
642 " new issue: https://github.com/vocdoni/arbo/issues/new") // TMP
643 case PrefixValueLeaf: // leaf
644 if !bytes.Equal(currValue, emptyValue) {
645 if getLeaf {
646 return currKey, currValue, siblings, nil
647 }
648 oldLeafKey, _ := ReadLeafValue(currValue)
649 if bytes.Equal(newKey, oldLeafKey) {
650 return nil, nil, nil, ErrKeyAlreadyExists
651 }
652
653 oldLeafKeyFull, err := keyPathFromKey(t.maxLevels, oldLeafKey)
654 if err != nil {
655 return nil, nil, nil, err
656 }
657
658 // if currKey is already used, go down until paths diverge
659 oldPath := getPath(t.maxLevels, oldLeafKeyFull)
660 siblings, err = t.downVirtually(siblings, currKey, newKey, oldPath, path, currLvl)
661 if err != nil {
662 return nil, nil, nil, err
663 }
664 }
665 return currKey, currValue, siblings, nil
666 case PrefixValueIntermediate: // intermediate
667 if len(currValue) != PrefixValueLen+t.hashFunction.Len()*2 {
668 return nil, nil, nil,
669 fmt.Errorf("intermediate value invalid length (expected: %d, actual: %d)",
670 PrefixValueLen+t.hashFunction.Len()*2, len(currValue))
671 }
672 // collect siblings while going down
673 if path[currLvl] {
674 // right
675 lChild, rChild := ReadIntermediateChilds(currValue)
676 siblings = append(siblings, lChild)
677 return t.down(rTx, newKey, rChild, siblings, path, currLvl+1, getLeaf)
678 }
679 // left
680 lChild, rChild := ReadIntermediateChilds(currValue)
681 siblings = append(siblings, rChild)
682 return t.down(rTx, newKey, lChild, siblings, path, currLvl+1, getLeaf)
683 default:
684 return nil, nil, nil, ErrInvalidValuePrefix
685 }
686}
687
688// downVirtually is used when in a leaf already exists, and a new leaf which
689// shares the path until the existing leaf is being added
690func (t *Tree) downVirtually(siblings [][]byte, oldKey, newKey []byte, oldPath,
691 newPath []bool, currLvl int) ([][]byte, error) {
692 var err error
693 if currLvl > t.maxLevels-1 {
694 return nil, ErrMaxVirtualLevel
695 }
696
697 if oldPath[currLvl] == newPath[currLvl] {
698 siblings = append(siblings, t.emptyHash)
699
700 siblings, err = t.downVirtually(siblings, oldKey, newKey, oldPath, newPath, currLvl+1)
701 if err != nil {
702 return nil, err
703 }
704 return siblings, nil
705 }
706 // reached the divergence
707 siblings = append(siblings, oldKey)
708
709 return siblings, nil
710}
711
712// up goes up recursively updating the intermediate nodes
713func (t *Tree) up(wTx db.WriteTx, key []byte, siblings [][]byte, path []bool,
714 currLvl, toLvl int) ([]byte, error) {
715 var k, v []byte
716 var err error
717 if path[currLvl+toLvl] {
718 k, v, err = t.newIntermediate(siblings[currLvl], key)
719 if err != nil {
720 return nil, err
721 }
722 } else {
723 k, v, err = t.newIntermediate(key, siblings[currLvl])
724 if err != nil {
725 return nil, err
726 }
727 }
728 // store k-v to db
729 if err = wTx.Set(k, v); err != nil {
730 return nil, err
731 }
732
733 if currLvl == 0 {
734 // reached the root
735 return k, nil
736 }
737
738 return t.up(wTx, k, siblings, path, currLvl-1, toLvl)
739}
740
741func (t *Tree) newLeafValue(k, v []byte) ([]byte, []byte, error) {
742 t.dbg.incHash()
743 return newLeafValue(t.hashFunction, k, v)
744}
745
746// newLeafValue takes a key & value from a leaf, and computes the leaf hash,
747// which is used as the leaf key. And the value is the concatenation of the
748// inputed key & value. The output of this function is used as key-value to
749// store the leaf in the DB.
750// [ 1 byte | 1 byte | N bytes | M bytes ]
751// [ type of node | length of key | key | value ]
752func newLeafValue(hashFunc HashFunction, k, v []byte) ([]byte, []byte, error) {
753 leafKey, err := hashFunc.Hash(k, v, []byte{1})
754 if err != nil {
755 return nil, nil, err
756 }
757 var leafValue []byte
758 leafValue = append(leafValue, byte(PrefixValueLeaf))
759 if len(k) > maxUint8 {
760 return nil, nil, fmt.Errorf("newLeafValue: len(k) > %v", maxUint8)
761 }
762 leafValue = append(leafValue, byte(len(k)))
763 leafValue = append(leafValue, k...)
764 leafValue = append(leafValue, v...)
765 return leafKey, leafValue, nil
766}
767
768// ReadLeafValue reads from a byte array the leaf key & value
769func ReadLeafValue(b []byte) ([]byte, []byte) {
770 if len(b) < PrefixValueLen {
771 return []byte{}, []byte{}
772 }
773
774 kLen := b[1]
775 if len(b) < PrefixValueLen+int(kLen) {
776 return []byte{}, []byte{}
777 }
778 k := b[PrefixValueLen : PrefixValueLen+kLen]
779 v := b[PrefixValueLen+kLen:]
780 return k, v
781}
782
783func (t *Tree) newIntermediate(l, r []byte) ([]byte, []byte, error) {
784 t.dbg.incHash()
785 return newIntermediate(t.hashFunction, l, r)
786}
787
788// newIntermediate takes the left & right keys of a intermediate node, and
789// computes its hash. Returns the hash of the node, which is the node key, and a
790// byte array that contains the value (which contains the left & right child
791// keys) to store in the DB.
792// [ 1 byte | 1 byte | N bytes | N bytes ]
793// [ type of node | length of left key | left key | right key ]
794func newIntermediate(hashFunc HashFunction, l, r []byte) ([]byte, []byte, error) {
795 b := make([]byte, PrefixValueLen+hashFunc.Len()*2)
796 b[0] = PrefixValueIntermediate
797 if len(l) > maxUint8 {
798 return nil, nil, fmt.Errorf("newIntermediate: len(l) > %v", maxUint8)
799 }
800 b[1] = byte(len(l))
801 copy(b[PrefixValueLen:PrefixValueLen+hashFunc.Len()], l)
802 copy(b[PrefixValueLen+hashFunc.Len():], r)
803
804 key, err := hashFunc.Hash(l, r)
805 if err != nil {
806 return nil, nil, err
807 }
808
809 return key, b, nil
810}
811
812// ReadIntermediateChilds reads from a byte array the two childs keys
813func ReadIntermediateChilds(b []byte) ([]byte, []byte) {
814 if len(b) < PrefixValueLen {
815 return []byte{}, []byte{}
816 }
817
818 lLen := b[1]
819 if len(b) < PrefixValueLen+int(lLen) {
820 return []byte{}, []byte{}
821 }
822 l := b[PrefixValueLen : PrefixValueLen+lLen]
823 r := b[PrefixValueLen+lLen:]
824 return l, r
825}
826
827func getPath(numLevels int, k []byte) []bool {
828 path := make([]bool, numLevels)
829 for n := 0; n < numLevels; n++ {
830 path[n] = k[n/8]&(1<<(n%8)) != 0
831 }
832 return path
833}
834
835// Update updates the value for a given existing key. If the given key does not
836// exist, returns an error.
837func (t *Tree) Update(k, v []byte) error {
838 wTx := t.db.WriteTx()
839 defer wTx.Discard()
840
841 if err := t.UpdateWithTx(wTx, k, v); err != nil {
842 return err
843 }
844 return wTx.Commit()
845}
846
847// UpdateWithTx does the same than the Update method, but allowing to pass the
848// db.WriteTx that is used. The db.WriteTx will not be committed inside this
849// method.
850func (t *Tree) UpdateWithTx(wTx db.WriteTx, k, v []byte) error {
851 t.Lock()
852 defer t.Unlock()
853
854 if !t.editable() {
855 return ErrSnapshotNotEditable
856 }
857
858 keyPath, err := keyPathFromKey(t.maxLevels, k)
859 if err != nil {
860 return err
861 }
862 path := getPath(t.maxLevels, keyPath)
863
864 root, err := t.RootWithTx(wTx)
865 if err != nil {
866 return err
867 }
868
869 var siblings [][]byte
870 _, valueAtBottom, siblings, err := t.down(wTx, k, root, siblings, path, 0, true)
871 if err != nil {
872 return err
873 }
874 oldKey, _ := ReadLeafValue(valueAtBottom)
875 if !bytes.Equal(oldKey, k) {
876 return ErrKeyNotFound
877 }
878
879 leafKey, leafValue, err := t.newLeafValue(k, v)
880 if err != nil {
881 return err
882 }
883
884 if err := wTx.Set(leafKey, leafValue); err != nil {
885 return err
886 }
887
888 // go up to the root
889 if len(siblings) == 0 {
890 return t.setRoot(wTx, leafKey)
891 }
892 root, err = t.up(wTx, leafKey, siblings, path, len(siblings)-1, 0)
893 if err != nil {
894 return err
895 }
896
897 // store root to db
898 if err := t.setRoot(wTx, root); err != nil {
899 return err
900 }
901 return nil
902}
903
904// GenProof generates a MerkleTree proof for the given key. The leaf value is
905// returned, together with the packed siblings of the proof, and a boolean
906// parameter that indicates if the proof is of existence (true) or not (false).
907func (t *Tree) GenProof(k []byte) ([]byte, []byte, []byte, bool, error) {
908 rTx := t.db.ReadTx()
909 defer rTx.Discard()
910
911 return t.GenProofWithTx(rTx, k)
912}
913
914// GenProofWithTx does the same than the GenProof method, but allowing to pass
915// the db.ReadTx that is used.
916func (t *Tree) GenProofWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, []byte, bool, error) {
917 keyPath, err := keyPathFromKey(t.maxLevels, k)
918 if err != nil {
919 return nil, nil, nil, false, err
920 }
921 path := getPath(t.maxLevels, keyPath)
922
923 root, err := t.RootWithTx(rTx)
924 if err != nil {
925 return nil, nil, nil, false, err
926 }
927
928 // go down to the leaf
929 var siblings [][]byte
930 _, value, siblings, err := t.down(rTx, k, root, siblings, path, 0, true)
931 if err != nil {
932 return nil, nil, nil, false, err
933 }
934
935 s, err := PackSiblings(t.hashFunction, siblings)
936 if err != nil {
937 return nil, nil, nil, false, err
938 }
939
940 leafK, leafV := ReadLeafValue(value)
941 if !bytes.Equal(k, leafK) {
942 // key not in tree, proof of non-existence
943 return leafK, leafV, s, false, nil
944 }
945
946 return leafK, leafV, s, true, nil
947}
948
949// PackSiblings packs the siblings into a byte array.
950// [ 2 byte | 2 byte | L bytes | S * N bytes ]
951// [ full length | bitmap length (L) | bitmap | N non-zero siblings ]
952// Where the bitmap indicates if the sibling is 0 or a value from the siblings
953// array. And S is the size of the output of the hash function used for the
954// Tree. The 2 2-byte that define the full length and bitmap length, are
955// encoded in little-endian.
956func PackSiblings(hashFunc HashFunction, siblings [][]byte) ([]byte, error) {
957 var b []byte
958 var bitmap []bool
959 emptySibling := make([]byte, hashFunc.Len())
960 for i := 0; i < len(siblings); i++ {
961 if bytes.Equal(siblings[i], emptySibling) {
962 bitmap = append(bitmap, false)
963 } else {
964 bitmap = append(bitmap, true)
965 b = append(b, siblings[i]...)
966 }
967 }
968
969 bitmapBytes := bitmapToBytes(bitmap)
970 l := len(bitmapBytes)
971 if l > maxUint16 {
972 return nil, fmt.Errorf("PackSiblings: bitmapBytes length > %v", maxUint16)
973 }
974
975 fullLen := 4 + l + len(b) //nolint:gomnd
976 if fullLen > maxUint16 {
977 return nil, fmt.Errorf("PackSiblings: fullLen > %v", maxUint16)
978 }
979 res := make([]byte, fullLen)
980 binary.LittleEndian.PutUint16(res[0:2], uint16(fullLen)) // set full length
981 binary.LittleEndian.PutUint16(res[2:4], uint16(l)) // set the bitmapBytes length
982 copy(res[4:4+l], bitmapBytes)
983 copy(res[4+l:], b)
984 return res, nil
985}
986
987// UnpackSiblings unpacks the siblings from a byte array.
988func UnpackSiblings(hashFunc HashFunction, b []byte) ([][]byte, error) {
989 fullLen := binary.LittleEndian.Uint16(b[0:2])
990 l := binary.LittleEndian.Uint16(b[2:4]) // bitmap bytes length
991 if len(b) != int(fullLen) {
992 return nil,
993 fmt.Errorf("expected len: %d, current len: %d",
994 fullLen, len(b))
995 }
996
997 bitmapBytes := b[4 : 4+l]
998 bitmap := bytesToBitmap(bitmapBytes)
999 siblingsBytes := b[4+l:]
1000 iSibl := 0
1001 emptySibl := make([]byte, hashFunc.Len())
1002 var siblings [][]byte
1003 for i := 0; i < len(bitmap); i++ {
1004 if iSibl >= len(siblingsBytes) {
1005 break
1006 }
1007 if bitmap[i] {
1008 siblings = append(siblings, siblingsBytes[iSibl:iSibl+hashFunc.Len()])
1009 iSibl += hashFunc.Len()
1010 } else {
1011 siblings = append(siblings, emptySibl)
1012 }
1013 }
1014 return siblings, nil
1015}
1016
1017func bitmapToBytes(bitmap []bool) []byte {
1018 bitmapBytesLen := int(math.Ceil(float64(len(bitmap)) / 8)) //nolint:gomnd
1019 b := make([]byte, bitmapBytesLen)
1020 for i := 0; i < len(bitmap); i++ {
1021 if bitmap[i] {
1022 b[i/8] |= 1 << (i % 8)
1023 }
1024 }
1025 return b
1026}
1027
1028func bytesToBitmap(b []byte) []bool {
1029 var bitmap []bool
1030 for i := 0; i < len(b); i++ {
1031 for j := 0; j < 8; j++ {
1032 bitmap = append(bitmap, b[i]&(1<<j) > 0)
1033 }
1034 }
1035 return bitmap
1036}
1037
1038// Get returns the value in the Tree for a given key. If the key is not found,
1039// will return the error ErrKeyNotFound, and in the leafK & leafV parameters
1040// will be placed the data found in the tree in the leaf that was on the path
1041// going to the input key.
1042func (t *Tree) Get(k []byte) ([]byte, []byte, error) {
1043 rTx := t.db.ReadTx()
1044 defer rTx.Discard()
1045
1046 return t.GetWithTx(rTx, k)
1047}
1048
1049// GetWithTx does the same than the Get method, but allowing to pass the
1050// db.ReadTx that is used. If the key is not found, will return the error
1051// ErrKeyNotFound, and in the leafK & leafV parameters will be placed the data
1052// found in the tree in the leaf that was on the path going to the input key.
1053func (t *Tree) GetWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, error) {
1054 keyPath, err := keyPathFromKey(t.maxLevels, k)
1055 if err != nil {
1056 return nil, nil, err
1057 }
1058 path := getPath(t.maxLevels, keyPath)
1059
1060 root, err := t.RootWithTx(rTx)
1061 if err != nil {
1062 return nil, nil, err
1063 }
1064
1065 // go down to the leaf
1066 var siblings [][]byte
1067 _, value, _, err := t.down(rTx, k, root, siblings, path, 0, true)
1068 if err != nil {
1069 return nil, nil, err
1070 }
1071 leafK, leafV := ReadLeafValue(value)
1072 if !bytes.Equal(k, leafK) {
1073 return leafK, leafV, ErrKeyNotFound
1074 }
1075
1076 return leafK, leafV, nil
1077}
1078
1079// CheckProof verifies the given proof. The proof verification depends on the
1080// HashFunction passed as parameter.
1081func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool, error) {
1082 siblings, err := UnpackSiblings(hashFunc, packedSiblings)
1083 if err != nil {
1084 return false, err
1085 }
1086
1087 keyPath := make([]byte, int(math.Ceil(float64(len(siblings))/float64(8)))) //nolint:gomnd
1088 copy(keyPath[:], k)
1089
1090 key, _, err := newLeafValue(hashFunc, k, v)
1091 if err != nil {
1092 return false, err
1093 }
1094
1095 path := getPath(len(siblings), keyPath)
1096 for i := len(siblings) - 1; i >= 0; i-- {
1097 if path[i] {
1098 key, _, err = newIntermediate(hashFunc, siblings[i], key)
1099 if err != nil {
1100 return false, err
1101 }
1102 } else {
1103 key, _, err = newIntermediate(hashFunc, key, siblings[i])
1104 if err != nil {
1105 return false, err
1106 }
1107 }
1108 }
1109 if bytes.Equal(key[:], root) {
1110 return true, nil
1111 }
1112 return false, nil
1113}
1114
1115func (t *Tree) incNLeafs(wTx db.WriteTx, nLeafs int) error {
1116 oldNLeafs, err := t.GetNLeafsWithTx(wTx)
1117 if err != nil {
1118 return err
1119 }
1120 newNLeafs := oldNLeafs + nLeafs
1121 return t.setNLeafs(wTx, newNLeafs)
1122}
1123
1124func (t *Tree) setNLeafs(wTx db.WriteTx, nLeafs int) error {
1125 b := make([]byte, 8)
1126 binary.LittleEndian.PutUint64(b, uint64(nLeafs))
1127 if err := wTx.Set(dbKeyNLeafs, b); err != nil {
1128 return err
1129 }
1130 return nil
1131}
1132
1133// GetNLeafs returns the number of Leafs of the Tree.
1134func (t *Tree) GetNLeafs() (int, error) {
1135 rTx := t.db.ReadTx()
1136 defer rTx.Discard()
1137
1138 return t.GetNLeafsWithTx(rTx)
1139}
1140
1141// GetNLeafsWithTx does the same than the GetNLeafs method, but allowing to
1142// pass the db.ReadTx that is used.
1143func (t *Tree) GetNLeafsWithTx(rTx db.ReadTx) (int, error) {
1144 b, err := rTx.Get(dbKeyNLeafs)
1145 if err != nil {
1146 return 0, err
1147 }
1148 nLeafs := binary.LittleEndian.Uint64(b)
1149 return int(nLeafs), nil
1150}
1151
1152// SetRoot sets the root to the given root
1153func (t *Tree) SetRoot(root []byte) error {
1154 wTx := t.db.WriteTx()
1155 defer wTx.Discard()
1156
1157 if err := t.SetRootWithTx(wTx, root); err != nil {
1158 return err
1159 }
1160 return wTx.Commit()
1161}
1162
1163// SetRootWithTx sets the root to the given root using the given db.WriteTx
1164func (t *Tree) SetRootWithTx(wTx db.WriteTx, root []byte) error {
1165 if !t.editable() {
1166 return ErrSnapshotNotEditable
1167 }
1168
1169 if root == nil {
1170 return fmt.Errorf("can not SetRoot with nil root")
1171 }
1172
1173 // check that the root exists in the db
1174 if !bytes.Equal(root, t.emptyHash) {
1175 if _, err := wTx.Get(root); err == ErrKeyNotFound {
1176 return fmt.Errorf("can not SetRoot with root %x, as it"+
1177 " does not exist in the db", root)
1178 } else if err != nil {
1179 return err
1180 }
1181 }
1182
1183 return wTx.Set(dbKeyRoot, root)
1184}
1185
1186// Snapshot returns a read-only copy of the Tree from the given root
1187func (t *Tree) Snapshot(fromRoot []byte) (*Tree, error) {
1188 // allow to define which root to use
1189 if fromRoot == nil {
1190 var err error
1191 fromRoot, err = t.Root()
1192 if err != nil {
1193 return nil, err
1194 }
1195 }
1196 rTx := t.db.ReadTx()
1197 defer rTx.Discard()
1198 // check that the root exists in the db
1199 if !bytes.Equal(fromRoot, t.emptyHash) {
1200 if _, err := rTx.Get(fromRoot); err == ErrKeyNotFound {
1201 return nil,
1202 fmt.Errorf("can not do a Snapshot with root %x,"+
1203 " as it does not exist in the db", fromRoot)
1204 } else if err != nil {
1205 return nil, err
1206 }
1207 }
1208
1209 return &Tree{
1210 db: t.db,
1211 maxLevels: t.maxLevels,
1212 snapshotRoot: fromRoot,
1213 emptyHash: t.emptyHash,
1214 hashFunction: t.hashFunction,
1215 dbg: t.dbg,
1216 }, nil
1217}
1218
1219// Iterate iterates through the full Tree, executing the given function on each
1220// node of the Tree.
1221func (t *Tree) Iterate(fromRoot []byte, f func([]byte, []byte)) error {
1222 rTx := t.db.ReadTx()
1223 defer rTx.Discard()
1224
1225 return t.IterateWithTx(rTx, fromRoot, f)
1226}
1227
1228// IterateWithTx does the same than the Iterate method, but allowing to pass
1229// the db.ReadTx that is used.
1230func (t *Tree) IterateWithTx(rTx db.ReadTx, fromRoot []byte, f func([]byte, []byte)) error {
1231 // allow to define which root to use
1232 if fromRoot == nil {
1233 var err error
1234 fromRoot, err = t.RootWithTx(rTx)
1235 if err != nil {
1236 return err
1237 }
1238 }
1239 return t.iter(rTx, fromRoot, f)
1240}
1241
1242// IterateWithStop does the same than Iterate, but with int for the current
1243// level, and a boolean parameter used by the passed function, is to indicate to
1244// stop iterating on the branch when the method returns 'true'.
1245func (t *Tree) IterateWithStop(fromRoot []byte, f func(int, []byte, []byte) bool) error {
1246 rTx := t.db.ReadTx()
1247 defer rTx.Discard()
1248
1249 // allow to define which root to use
1250 if fromRoot == nil {
1251 var err error
1252 fromRoot, err = t.RootWithTx(rTx)
1253 if err != nil {
1254 return err
1255 }
1256 }
1257 return t.iterWithStop(rTx, fromRoot, 0, f)
1258}
1259
1260// IterateWithStopWithTx does the same than the IterateWithStop method, but
1261// allowing to pass the db.ReadTx that is used.
1262func (t *Tree) IterateWithStopWithTx(rTx db.ReadTx, fromRoot []byte,
1263 f func(int, []byte, []byte) bool) error {
1264 // allow to define which root to use
1265 if fromRoot == nil {
1266 var err error
1267 fromRoot, err = t.RootWithTx(rTx)
1268 if err != nil {
1269 return err
1270 }
1271 }
1272 return t.iterWithStop(rTx, fromRoot, 0, f)
1273}
1274
1275func (t *Tree) iterWithStop(rTx db.ReadTx, k []byte, currLevel int,
1276 f func(int, []byte, []byte) bool) error {
1277 var v []byte
1278 var err error
1279 if bytes.Equal(k, t.emptyHash) {
1280 v = t.emptyHash
1281 } else {
1282 v, err = rTx.Get(k)
1283 if err != nil {
1284 return err
1285 }
1286 }
1287 currLevel++
1288
1289 switch v[0] {
1290 case PrefixValueEmpty:
1291 f(currLevel, k, v)
1292 case PrefixValueLeaf:
1293 f(currLevel, k, v)
1294 case PrefixValueIntermediate:
1295 stop := f(currLevel, k, v)
1296 if stop {
1297 return nil
1298 }
1299 l, r := ReadIntermediateChilds(v)
1300 if err = t.iterWithStop(rTx, l, currLevel, f); err != nil {
1301 return err
1302 }
1303 if err = t.iterWithStop(rTx, r, currLevel, f); err != nil {
1304 return err
1305 }
1306 default:
1307 return ErrInvalidValuePrefix
1308 }
1309 return nil
1310}
1311
1312func (t *Tree) iter(rTx db.ReadTx, k []byte, f func([]byte, []byte)) error {
1313 f2 := func(currLvl int, k, v []byte) bool {
1314 f(k, v)
1315 return false
1316 }
1317 return t.iterWithStop(rTx, k, 0, f2)
1318}
1319
1320// Dump exports all the Tree leafs in a byte array of length:
1321// [ N * (2+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
1322// [ 1 byte | 1 byte | S bytes | len(v) bytes ]
1323// [ len(k) | len(v) | key | value ]
1324// Where S is the size of the output of the hash function used for the Tree.
1325func (t *Tree) Dump(fromRoot []byte) ([]byte, error) {
1326 // allow to define which root to use
1327 if fromRoot == nil {
1328 var err error
1329 fromRoot, err = t.Root()
1330 if err != nil {
1331 return nil, err
1332 }
1333 }
1334
1335 // WARNING current encoding only supports key & values of 255 bytes each
1336 // (due using only 1 byte for the length headers).
1337 var b []byte
1338 var callbackErr error
1339 err := t.IterateWithStop(fromRoot, func(_ int, k, v []byte) bool {
1340 if v[0] != PrefixValueLeaf {
1341 return false
1342 }
1343 leafK, leafV := ReadLeafValue(v)
1344 kv := make([]byte, 2+len(leafK)+len(leafV))
1345 if len(leafK) > maxUint8 {
1346 callbackErr = fmt.Errorf("len(leafK) > %v", maxUint8)
1347 return true
1348 }
1349 kv[0] = byte(len(leafK))
1350 if len(leafV) > maxUint8 {
1351 callbackErr = fmt.Errorf("len(leafV) > %v", maxUint8)
1352 return true
1353 }
1354 kv[1] = byte(len(leafV))
1355 copy(kv[2:2+len(leafK)], leafK)
1356 copy(kv[2+len(leafK):], leafV)
1357 b = append(b, kv...)
1358 return false
1359 })
1360 if callbackErr != nil {
1361 return nil, callbackErr
1362 }
1363 return b, err
1364}
1365
1366// ImportDump imports the leafs (that have been exported with the Dump method)
1367// in the Tree.
1368func (t *Tree) ImportDump(b []byte) error {
1369 if !t.editable() {
1370 return ErrSnapshotNotEditable
1371 }
1372 root, err := t.Root()
1373 if err != nil {
1374 return err
1375 }
1376 if !bytes.Equal(root, t.emptyHash) {
1377 return ErrTreeNotEmpty
1378 }
1379
1380 r := bytes.NewReader(b)
1381 var keys, values [][]byte
1382 for {
1383 l := make([]byte, 2)
1384 _, err = io.ReadFull(r, l)
1385 if err == io.EOF {
1386 break
1387 } else if err != nil {
1388 return err
1389 }
1390 k := make([]byte, l[0])
1391 _, err = io.ReadFull(r, k)
1392 if err != nil {
1393 return err
1394 }
1395 v := make([]byte, l[1])
1396 _, err = io.ReadFull(r, v)
1397 if err != nil {
1398 return err
1399 }
1400 keys = append(keys, k)
1401 values = append(values, v)
1402 }
1403 if _, err = t.AddBatch(keys, values); err != nil {
1404 return err
1405 }
1406 return nil
1407}
1408
1409// Graphviz iterates across the full tree to generate a string Graphviz
1410// representation of the tree and writes it to w
1411func (t *Tree) Graphviz(w io.Writer, fromRoot []byte) error {
1412 return t.GraphvizFirstNLevels(w, fromRoot, t.maxLevels)
1413}
1414
1415// GraphvizFirstNLevels iterates across the first NLevels of the tree to
1416// generate a string Graphviz representation of the first NLevels of the tree
1417// and writes it to w
1418func (t *Tree) GraphvizFirstNLevels(w io.Writer, fromRoot []byte, untilLvl int) error {
1419 fmt.Fprintf(w, `digraph hierarchy {
1420node [fontname=Monospace,fontsize=10,shape=box]
1421`)
1422
1423 rTx := t.db.ReadTx()
1424 defer rTx.Discard()
1425
1426 if fromRoot == nil {
1427 var err error
1428 fromRoot, err = t.RootWithTx(rTx)
1429 if err != nil {
1430 return err
1431 }
1432 }
1433
1434 nEmpties := 0
1435 err := t.iterWithStop(rTx, fromRoot, 0, func(currLvl int, k, v []byte) bool {
1436 if currLvl == untilLvl {
1437 return true // to stop the iter from going down
1438 }
1439 switch v[0] {
1440 case PrefixValueEmpty:
1441 case PrefixValueLeaf:
1442 fmt.Fprintf(w, "\"%v\" [style=filled];\n", hex.EncodeToString(k[:nChars]))
1443 // key & value from the leaf
1444 kB, vB := ReadLeafValue(v)
1445 fmt.Fprintf(w, "\"%v\" -> {\"k:%v\\nv:%v\"}\n",
1446 hex.EncodeToString(k[:nChars]), hex.EncodeToString(kB[:nChars]),
1447 hex.EncodeToString(vB[:nChars]))
1448 fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
1449 hex.EncodeToString(kB[:nChars]), hex.EncodeToString(vB[:nChars]))
1450 case PrefixValueIntermediate:
1451 l, r := ReadIntermediateChilds(v)
1452 lStr := hex.EncodeToString(l[:nChars])
1453 rStr := hex.EncodeToString(r[:nChars])
1454 eStr := ""
1455 if bytes.Equal(l, t.emptyHash) {
1456 lStr = fmt.Sprintf("empty%v", nEmpties)
1457 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
1458 lStr)
1459 nEmpties++
1460 }
1461 if bytes.Equal(r, t.emptyHash) {
1462 rStr = fmt.Sprintf("empty%v", nEmpties)
1463 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
1464 rStr)
1465 nEmpties++
1466 }
1467 fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", hex.EncodeToString(k[:nChars]),
1468 lStr, rStr)
1469 fmt.Fprint(w, eStr)
1470 default:
1471 }
1472 return false
1473 })
1474 fmt.Fprintf(w, "}\n")
1475 return err
1476}
1477
1478// PrintGraphviz prints the output of Tree.Graphviz
1479func (t *Tree) PrintGraphviz(fromRoot []byte) error {
1480 if fromRoot == nil {
1481 var err error
1482 fromRoot, err = t.Root()
1483 if err != nil {
1484 return err
1485 }
1486 }
1487 return t.PrintGraphvizFirstNLevels(fromRoot, t.maxLevels)
1488}
1489
1490// PrintGraphvizFirstNLevels prints the output of Tree.GraphvizFirstNLevels
1491func (t *Tree) PrintGraphvizFirstNLevels(fromRoot []byte, untilLvl int) error {
1492 if fromRoot == nil {
1493 var err error
1494 fromRoot, err = t.Root()
1495 if err != nil {
1496 return err
1497 }
1498 }
1499 w := bytes.NewBufferString("")
1500 fmt.Fprintf(w,
1501 "--------\nGraphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+":\n")
1502 err := t.GraphvizFirstNLevels(w, fromRoot, untilLvl)
1503 if err != nil {
1504 fmt.Println(w)
1505 return err
1506 }
1507 fmt.Fprintf(w,
1508 "End of Graphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+"\n--------\n")
1509
1510 fmt.Println(w)
1511 return nil
1512}
1513
1514// TODO circom proofs
1515// TODO data structure for proofs (including root, key, value, siblings,
1516// hashFunction) + method to verify that data structure