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
380 var newSubRoots [][]byte
381 // TODO store the nodes key-values into the wTx
382 for i := 0; i < len(subRoots); i += 2 {
383 if bytes.Equal(subRoots[i], t.emptyHash) && bytes.Equal(subRoots[i+1], t.emptyHash) {
384 // TODO: || (ns[i].typ() == vtLeaf && ns[i+1].typ() == vtEmpty) {
385
386 // when both sub nodes are empty, the parent is also empty
387 // or
388 // (TODO WIP) when 1st sub node is a leaf but the 2nd is empty, the
389 // leaf is used as parent
390
391 newSubRoots = append(newSubRoots, subRoots[i])
392 continue
393 }
394 // TODO if ns[i].typ() == vtEmpty && ns[i+1].typ() == vtLeaf {
395 // when 2nd sub node is a leaf but the 1st is empty, the
396 // leaf is used as 'parent'
397
398 k, v, err := t.newIntermediate(subRoots[i], subRoots[i+1])
399 if err != nil {
400 return nil, err
401 }
402 // store k-v to db
403 if err = wTx.Set(k, v); err != nil {
404 return nil, err
405 }
406 newSubRoots = append(newSubRoots, k)
407 }
408
409 return t.upFromSubRoots(wTx, newSubRoots)
410}
411func (t *Tree) getSubRootsAtLevel(rTx db.ReadTx, root []byte, l int) ([][]byte, error) {
412 // go at level l and return each node key, where each node key is the
413 // subRoot of the subTree that starts there
414
415 var subRoots [][]byte
416 err := t.iterWithStop(rTx, root, 0, func(currLvl int, k, v []byte) bool {
417 if currLvl == l && !bytes.Equal(k, t.emptyHash) {
418 subRoots = append(subRoots, k)
419 }
420 if currLvl >= l {
421 return true // to stop the iter from going down
422 }
423 return false
424 })
425
426 return subRoots, err
427}
428
429func (t *Tree) addBatchInMemory(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
430 vt, err := t.loadVT(wTx)
431 if err != nil {
432 return nil, err
433 }
434
435 invalids, err := vt.addBatch(keys, values)
436 if err != nil {
437 return nil, err
438 }
439
440 // once the VirtualTree is build, compute the hashes
441 pairs, err := vt.computeHashes()
442 if err != nil {
443 // currently invalids in computeHashes are not counted,
444 // but should not be needed, as if there is an error there is
445 // nothing stored in the db and the error is returned
446 return nil, err
447 }
448
449 // store pairs in db
450 for i := 0; i < len(pairs); i++ {
451 if err := wTx.Set(pairs[i][0], pairs[i][1]); err != nil {
452 return nil, err
453 }
454 }
455
456 // store root (from the vt) to db
457 if vt.root != nil {
458 if err := wTx.Set(dbKeyRoot, vt.root.h); err != nil {
459 return nil, err
460 }
461 }
462
463 // update nLeafs
464 if err := t.incNLeafs(wTx, len(keys)-len(invalids)); err != nil {
465 return nil, err
466 }
467
468 return invalids, nil
469}
470
471// loadVT loads a new virtual tree (vt) from the current Tree, which contains
472// the same leafs.
473func (t *Tree) loadVT(rTx db.ReadTx) (vt, error) {
474 vt := newVT(t.maxLevels, t.hashFunction)
475 vt.params.dbg = t.dbg
476 var callbackErr error
477 err := t.IterateWithStopWithTx(rTx, nil, func(_ int, k, v []byte) bool {
478 if v[0] != PrefixValueLeaf {
479 return false
480 }
481 leafK, leafV := ReadLeafValue(v)
482 if err := vt.add(0, leafK, leafV); err != nil {
483 callbackErr = err
484 return true
485 }
486 return false
487 })
488 if callbackErr != nil {
489 return vt, callbackErr
490 }
491
492 return vt, err
493}
494
495// Add inserts the key-value into the Tree. If the inputs come from a
496// *big.Int, is expected that are represented by a Little-Endian byte array
497// (for circom compatibility).
498func (t *Tree) Add(k, v []byte) error {
499 wTx := t.db.WriteTx()
500 defer wTx.Discard()
501
502 if err := t.AddWithTx(wTx, k, v); err != nil {
503 return err
504 }
505
506 return wTx.Commit()
507}
508
509// AddWithTx does the same than the Add method, but allowing to pass the
510// db.WriteTx that is used. The db.WriteTx will not be committed inside this
511// method.
512func (t *Tree) AddWithTx(wTx db.WriteTx, k, v []byte) error {
513 t.Lock()
514 defer t.Unlock()
515
516 if !t.editable() {
517 return ErrSnapshotNotEditable
518 }
519 return t.addWithTx(wTx, k, v)
520}
521
522// warning: addWithTx does not use the Tree mutex, the mutex is responsibility
523// of the methods calling this method, and same with t.editable().
524func (t *Tree) addWithTx(wTx db.WriteTx, k, v []byte) error {
525 root, err := t.RootWithTx(wTx)
526 if err != nil {
527 return err
528 }
529
530 root, err = t.add(wTx, root, 0, k, v) // add from level 0
531 if err != nil {
532 return err
533 }
534 // store root to db
535 if err := t.setRoot(wTx, root); err != nil {
536 return err
537 }
538 // update nLeafs
539 if err = t.incNLeafs(wTx, 1); err != nil {
540 return err
541 }
542 return nil
543}
544
545// keyPathFromKey returns the keyPath and checks that the key is not bigger
546// than maximum key length for the tree maxLevels size.
547// This is because if the key bits length is bigger than the maxLevels of the
548// tree, two different keys that their difference is at the end, will collision
549// in the same leaf of the tree (at the max depth).
550func keyPathFromKey(maxLevels int, k []byte) ([]byte, error) {
551 maxKeyLen := int(math.Ceil(float64(maxLevels) / float64(8))) //nolint:gomnd
552 if len(k) > maxKeyLen {
553 return nil, fmt.Errorf("len(k) can not be bigger than ceil(maxLevels/8), where"+
554 " len(k): %d, maxLevels: %d, max key len=ceil(maxLevels/8): %d. Might need"+
555 " a bigger tree depth (maxLevels>=%d) in order to input keys of length %d",
556 len(k), maxLevels, maxKeyLen, len(k)*8, len(k)) //nolint:gomnd
557 }
558 keyPath := make([]byte, maxKeyLen) //nolint:gomnd
559 copy(keyPath[:], k)
560 return keyPath, nil
561}
562
563func (t *Tree) add(wTx db.WriteTx, root []byte, fromLvl int, k, v []byte) ([]byte, error) {
564 keyPath, err := keyPathFromKey(t.maxLevels, k)
565 if err != nil {
566 return nil, err
567 }
568 path := getPath(t.maxLevels, keyPath)
569
570 // go down to the leaf
571 var siblings [][]byte
572 _, _, siblings, err = t.down(wTx, k, root, siblings, path, fromLvl, false)
573 if err != nil {
574 return nil, err
575 }
576
577 leafKey, leafValue, err := t.newLeafValue(k, v)
578 if err != nil {
579 return nil, err
580 }
581
582 if err := wTx.Set(leafKey, leafValue); err != nil {
583 return nil, err
584 }
585
586 // go up to the root
587 if len(siblings) == 0 {
588 // return the leafKey as root
589 return leafKey, nil
590 }
591 root, err = t.up(wTx, leafKey, siblings, path, len(siblings)-1, fromLvl)
592 if err != nil {
593 return nil, err
594 }
595
596 return root, nil
597}
598
599// down goes down to the leaf recursively
600func (t *Tree) down(rTx db.ReadTx, newKey, currKey []byte, siblings [][]byte,
601 path []bool, currLvl int, getLeaf bool) (
602 []byte, []byte, [][]byte, error) {
603 if currLvl > t.maxLevels {
604 return nil, nil, nil, ErrMaxLevel
605 }
606
607 var err error
608 var currValue []byte
609 if bytes.Equal(currKey, t.emptyHash) {
610 // empty value
611 return currKey, emptyValue, siblings, nil
612 }
613 currValue, err = rTx.Get(currKey)
614 if err != nil {
615 return nil, nil, nil, err
616 }
617
618 switch currValue[0] {
619 case PrefixValueEmpty: // empty
620 fmt.Printf("newKey: %s, currKey: %s, currLvl: %d, currValue: %s\n",
621 hex.EncodeToString(newKey), hex.EncodeToString(currKey),
622 currLvl, hex.EncodeToString(currValue))
623 panic("This point should not be reached, as the 'if currKey==t.emptyHash'" +
624 " above should avoid reaching this point. This panic is temporary" +
625 " for reporting purposes, will be deleted in future versions." +
626 " Please paste this log (including the previous log lines) in a" +
627 " new issue: https://github.com/vocdoni/arbo/issues/new") // TMP
628 case PrefixValueLeaf: // leaf
629 if !bytes.Equal(currValue, emptyValue) {
630 if getLeaf {
631 return currKey, currValue, siblings, nil
632 }
633 oldLeafKey, _ := ReadLeafValue(currValue)
634 if bytes.Equal(newKey, oldLeafKey) {
635 return nil, nil, nil, ErrKeyAlreadyExists
636 }
637
638 oldLeafKeyFull, err := keyPathFromKey(t.maxLevels, oldLeafKey)
639 if err != nil {
640 return nil, nil, nil, err
641 }
642
643 // if currKey is already used, go down until paths diverge
644 oldPath := getPath(t.maxLevels, oldLeafKeyFull)
645 siblings, err = t.downVirtually(siblings, currKey, newKey, oldPath, path, currLvl)
646 if err != nil {
647 return nil, nil, nil, err
648 }
649 }
650 return currKey, currValue, siblings, nil
651 case PrefixValueIntermediate: // intermediate
652 if len(currValue) != PrefixValueLen+t.hashFunction.Len()*2 {
653 return nil, nil, nil,
654 fmt.Errorf("intermediate value invalid length (expected: %d, actual: %d)",
655 PrefixValueLen+t.hashFunction.Len()*2, len(currValue))
656 }
657 // collect siblings while going down
658 if path[currLvl] {
659 // right
660 lChild, rChild := ReadIntermediateChilds(currValue)
661 siblings = append(siblings, lChild)
662 return t.down(rTx, newKey, rChild, siblings, path, currLvl+1, getLeaf)
663 }
664 // left
665 lChild, rChild := ReadIntermediateChilds(currValue)
666 siblings = append(siblings, rChild)
667 return t.down(rTx, newKey, lChild, siblings, path, currLvl+1, getLeaf)
668 default:
669 return nil, nil, nil, ErrInvalidValuePrefix
670 }
671}
672
673// downVirtually is used when in a leaf already exists, and a new leaf which
674// shares the path until the existing leaf is being added
675func (t *Tree) downVirtually(siblings [][]byte, oldKey, newKey []byte, oldPath,
676 newPath []bool, currLvl int) ([][]byte, error) {
677 var err error
678 if currLvl > t.maxLevels-1 {
679 return nil, ErrMaxVirtualLevel
680 }
681
682 if oldPath[currLvl] == newPath[currLvl] {
683 siblings = append(siblings, t.emptyHash)
684
685 siblings, err = t.downVirtually(siblings, oldKey, newKey, oldPath, newPath, currLvl+1)
686 if err != nil {
687 return nil, err
688 }
689 return siblings, nil
690 }
691 // reached the divergence
692 siblings = append(siblings, oldKey)
693
694 return siblings, nil
695}
696
697// up goes up recursively updating the intermediate nodes
698func (t *Tree) up(wTx db.WriteTx, key []byte, siblings [][]byte, path []bool,
699 currLvl, toLvl int) ([]byte, error) {
700 var k, v []byte
701 var err error
702 if path[currLvl+toLvl] {
703 k, v, err = t.newIntermediate(siblings[currLvl], key)
704 if err != nil {
705 return nil, err
706 }
707 } else {
708 k, v, err = t.newIntermediate(key, siblings[currLvl])
709 if err != nil {
710 return nil, err
711 }
712 }
713 // store k-v to db
714 if err = wTx.Set(k, v); err != nil {
715 return nil, err
716 }
717
718 if currLvl == 0 {
719 // reached the root
720 return k, nil
721 }
722
723 return t.up(wTx, k, siblings, path, currLvl-1, toLvl)
724}
725
726func (t *Tree) newLeafValue(k, v []byte) ([]byte, []byte, error) {
727 t.dbg.incHash()
728 return newLeafValue(t.hashFunction, k, v)
729}
730
731// newLeafValue takes a key & value from a leaf, and computes the leaf hash,
732// which is used as the leaf key. And the value is the concatenation of the
733// inputed key & value. The output of this function is used as key-value to
734// store the leaf in the DB.
735// [ 1 byte | 1 byte | N bytes | M bytes ]
736// [ type of node | length of key | key | value ]
737func newLeafValue(hashFunc HashFunction, k, v []byte) ([]byte, []byte, error) {
738 leafKey, err := hashFunc.Hash(k, v, []byte{1})
739 if err != nil {
740 return nil, nil, err
741 }
742 var leafValue []byte
743 leafValue = append(leafValue, byte(PrefixValueLeaf))
744 if len(k) > maxUint8 {
745 return nil, nil, fmt.Errorf("newLeafValue: len(k) > %v", maxUint8)
746 }
747 leafValue = append(leafValue, byte(len(k)))
748 leafValue = append(leafValue, k...)
749 leafValue = append(leafValue, v...)
750 return leafKey, leafValue, nil
751}
752
753// ReadLeafValue reads from a byte array the leaf key & value
754func ReadLeafValue(b []byte) ([]byte, []byte) {
755 if len(b) < PrefixValueLen {
756 return []byte{}, []byte{}
757 }
758
759 kLen := b[1]
760 if len(b) < PrefixValueLen+int(kLen) {
761 return []byte{}, []byte{}
762 }
763 k := b[PrefixValueLen : PrefixValueLen+kLen]
764 v := b[PrefixValueLen+kLen:]
765 return k, v
766}
767
768func (t *Tree) newIntermediate(l, r []byte) ([]byte, []byte, error) {
769 t.dbg.incHash()
770 return newIntermediate(t.hashFunction, l, r)
771}
772
773// newIntermediate takes the left & right keys of a intermediate node, and
774// computes its hash. Returns the hash of the node, which is the node key, and a
775// byte array that contains the value (which contains the left & right child
776// keys) to store in the DB.
777// [ 1 byte | 1 byte | N bytes | N bytes ]
778// [ type of node | length of left key | left key | right key ]
779func newIntermediate(hashFunc HashFunction, l, r []byte) ([]byte, []byte, error) {
780 b := make([]byte, PrefixValueLen+hashFunc.Len()*2)
781 b[0] = PrefixValueIntermediate
782 if len(l) > maxUint8 {
783 return nil, nil, fmt.Errorf("newIntermediate: len(l) > %v", maxUint8)
784 }
785 b[1] = byte(len(l))
786 copy(b[PrefixValueLen:PrefixValueLen+hashFunc.Len()], l)
787 copy(b[PrefixValueLen+hashFunc.Len():], r)
788
789 key, err := hashFunc.Hash(l, r)
790 if err != nil {
791 return nil, nil, err
792 }
793
794 return key, b, nil
795}
796
797// ReadIntermediateChilds reads from a byte array the two childs keys
798func ReadIntermediateChilds(b []byte) ([]byte, []byte) {
799 if len(b) < PrefixValueLen {
800 return []byte{}, []byte{}
801 }
802
803 lLen := b[1]
804 if len(b) < PrefixValueLen+int(lLen) {
805 return []byte{}, []byte{}
806 }
807 l := b[PrefixValueLen : PrefixValueLen+lLen]
808 r := b[PrefixValueLen+lLen:]
809 return l, r
810}
811
812func getPath(numLevels int, k []byte) []bool {
813 path := make([]bool, numLevels)
814 for n := 0; n < numLevels; n++ {
815 path[n] = k[n/8]&(1<<(n%8)) != 0
816 }
817 return path
818}
819
820// Update updates the value for a given existing key. If the given key does not
821// exist, returns an error.
822func (t *Tree) Update(k, v []byte) error {
823 wTx := t.db.WriteTx()
824 defer wTx.Discard()
825
826 if err := t.UpdateWithTx(wTx, k, v); err != nil {
827 return err
828 }
829 return wTx.Commit()
830}
831
832// UpdateWithTx does the same than the Update method, but allowing to pass the
833// db.WriteTx that is used. The db.WriteTx will not be committed inside this
834// method.
835func (t *Tree) UpdateWithTx(wTx db.WriteTx, k, v []byte) error {
836 t.Lock()
837 defer t.Unlock()
838
839 if !t.editable() {
840 return ErrSnapshotNotEditable
841 }
842
843 keyPath, err := keyPathFromKey(t.maxLevels, k)
844 if err != nil {
845 return err
846 }
847 path := getPath(t.maxLevels, keyPath)
848
849 root, err := t.RootWithTx(wTx)
850 if err != nil {
851 return err
852 }
853
854 var siblings [][]byte
855 _, valueAtBottom, siblings, err := t.down(wTx, k, root, siblings, path, 0, true)
856 if err != nil {
857 return err
858 }
859 oldKey, _ := ReadLeafValue(valueAtBottom)
860 if !bytes.Equal(oldKey, k) {
861 return ErrKeyNotFound
862 }
863
864 leafKey, leafValue, err := t.newLeafValue(k, v)
865 if err != nil {
866 return err
867 }
868
869 if err := wTx.Set(leafKey, leafValue); err != nil {
870 return err
871 }
872
873 // go up to the root
874 if len(siblings) == 0 {
875 return t.setRoot(wTx, leafKey)
876 }
877 root, err = t.up(wTx, leafKey, siblings, path, len(siblings)-1, 0)
878 if err != nil {
879 return err
880 }
881
882 // store root to db
883 if err := t.setRoot(wTx, root); err != nil {
884 return err
885 }
886 return nil
887}
888
889// GenProof generates a MerkleTree proof for the given key. The leaf value is
890// returned, together with the packed siblings of the proof, and a boolean
891// parameter that indicates if the proof is of existence (true) or not (false).
892func (t *Tree) GenProof(k []byte) ([]byte, []byte, []byte, bool, error) {
893 rTx := t.db.ReadTx()
894 defer rTx.Discard()
895
896 return t.GenProofWithTx(rTx, k)
897}
898
899// GenProofWithTx does the same than the GenProof method, but allowing to pass
900// the db.ReadTx that is used.
901func (t *Tree) GenProofWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, []byte, bool, error) {
902 keyPath, err := keyPathFromKey(t.maxLevels, k)
903 if err != nil {
904 return nil, nil, nil, false, err
905 }
906 path := getPath(t.maxLevels, keyPath)
907
908 root, err := t.RootWithTx(rTx)
909 if err != nil {
910 return nil, nil, nil, false, err
911 }
912
913 // go down to the leaf
914 var siblings [][]byte
915 _, value, siblings, err := t.down(rTx, k, root, siblings, path, 0, true)
916 if err != nil {
917 return nil, nil, nil, false, err
918 }
919
920 s, err := PackSiblings(t.hashFunction, siblings)
921 if err != nil {
922 return nil, nil, nil, false, err
923 }
924
925 leafK, leafV := ReadLeafValue(value)
926 if !bytes.Equal(k, leafK) {
927 // key not in tree, proof of non-existence
928 return leafK, leafV, s, false, nil
929 }
930
931 return leafK, leafV, s, true, nil
932}
933
934// PackSiblings packs the siblings into a byte array.
935// [ 2 byte | 2 byte | L bytes | S * N bytes ]
936// [ full length | bitmap length (L) | bitmap | N non-zero siblings ]
937// Where the bitmap indicates if the sibling is 0 or a value from the siblings
938// array. And S is the size of the output of the hash function used for the
939// Tree. The 2 2-byte that define the full length and bitmap length, are
940// encoded in little-endian.
941func PackSiblings(hashFunc HashFunction, siblings [][]byte) ([]byte, error) {
942 var b []byte
943 var bitmap []bool
944 emptySibling := make([]byte, hashFunc.Len())
945 for i := 0; i < len(siblings); i++ {
946 if bytes.Equal(siblings[i], emptySibling) {
947 bitmap = append(bitmap, false)
948 } else {
949 bitmap = append(bitmap, true)
950 b = append(b, siblings[i]...)
951 }
952 }
953
954 bitmapBytes := bitmapToBytes(bitmap)
955 l := len(bitmapBytes)
956 if l > maxUint16 {
957 return nil, fmt.Errorf("PackSiblings: bitmapBytes length > %v", maxUint16)
958 }
959
960 fullLen := 4 + l + len(b) //nolint:gomnd
961 if fullLen > maxUint16 {
962 return nil, fmt.Errorf("PackSiblings: fullLen > %v", maxUint16)
963 }
964 res := make([]byte, fullLen)
965 binary.LittleEndian.PutUint16(res[0:2], uint16(fullLen)) // set full length
966 binary.LittleEndian.PutUint16(res[2:4], uint16(l)) // set the bitmapBytes length
967 copy(res[4:4+l], bitmapBytes)
968 copy(res[4+l:], b)
969 return res, nil
970}
971
972// UnpackSiblings unpacks the siblings from a byte array.
973func UnpackSiblings(hashFunc HashFunction, b []byte) ([][]byte, error) {
974 fullLen := binary.LittleEndian.Uint16(b[0:2])
975 l := binary.LittleEndian.Uint16(b[2:4]) // bitmap bytes length
976 if len(b) != int(fullLen) {
977 return nil,
978 fmt.Errorf("expected len: %d, current len: %d",
979 fullLen, len(b))
980 }
981
982 bitmapBytes := b[4 : 4+l]
983 bitmap := bytesToBitmap(bitmapBytes)
984 siblingsBytes := b[4+l:]
985 iSibl := 0
986 emptySibl := make([]byte, hashFunc.Len())
987 var siblings [][]byte
988 for i := 0; i < len(bitmap); i++ {
989 if iSibl >= len(siblingsBytes) {
990 break
991 }
992 if bitmap[i] {
993 siblings = append(siblings, siblingsBytes[iSibl:iSibl+hashFunc.Len()])
994 iSibl += hashFunc.Len()
995 } else {
996 siblings = append(siblings, emptySibl)
997 }
998 }
999 return siblings, nil
1000}
1001
1002func bitmapToBytes(bitmap []bool) []byte {
1003 bitmapBytesLen := int(math.Ceil(float64(len(bitmap)) / 8)) //nolint:gomnd
1004 b := make([]byte, bitmapBytesLen)
1005 for i := 0; i < len(bitmap); i++ {
1006 if bitmap[i] {
1007 b[i/8] |= 1 << (i % 8)
1008 }
1009 }
1010 return b
1011}
1012
1013func bytesToBitmap(b []byte) []bool {
1014 var bitmap []bool
1015 for i := 0; i < len(b); i++ {
1016 for j := 0; j < 8; j++ {
1017 bitmap = append(bitmap, b[i]&(1<<j) > 0)
1018 }
1019 }
1020 return bitmap
1021}
1022
1023// Get returns the value in the Tree for a given key. If the key is not found,
1024// will return the error ErrKeyNotFound, and in the leafK & leafV parameters
1025// will be placed the data found in the tree in the leaf that was on the path
1026// going to the input key.
1027func (t *Tree) Get(k []byte) ([]byte, []byte, error) {
1028 rTx := t.db.ReadTx()
1029 defer rTx.Discard()
1030
1031 return t.GetWithTx(rTx, k)
1032}
1033
1034// GetWithTx does the same than the Get method, but allowing to pass the
1035// db.ReadTx that is used. If the key is not found, will return the error
1036// ErrKeyNotFound, and in the leafK & leafV parameters will be placed the data
1037// found in the tree in the leaf that was on the path going to the input key.
1038func (t *Tree) GetWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, error) {
1039 keyPath, err := keyPathFromKey(t.maxLevels, k)
1040 if err != nil {
1041 return nil, nil, err
1042 }
1043 path := getPath(t.maxLevels, keyPath)
1044
1045 root, err := t.RootWithTx(rTx)
1046 if err != nil {
1047 return nil, nil, err
1048 }
1049
1050 // go down to the leaf
1051 var siblings [][]byte
1052 _, value, _, err := t.down(rTx, k, root, siblings, path, 0, true)
1053 if err != nil {
1054 return nil, nil, err
1055 }
1056 leafK, leafV := ReadLeafValue(value)
1057 if !bytes.Equal(k, leafK) {
1058 return leafK, leafV, ErrKeyNotFound
1059 }
1060
1061 return leafK, leafV, nil
1062}
1063
1064// CheckProof verifies the given proof. The proof verification depends on the
1065// HashFunction passed as parameter.
1066func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool, error) {
1067 siblings, err := UnpackSiblings(hashFunc, packedSiblings)
1068 if err != nil {
1069 return false, err
1070 }
1071
1072 keyPath := make([]byte, int(math.Ceil(float64(len(siblings))/float64(8)))) //nolint:gomnd
1073 copy(keyPath[:], k)
1074
1075 key, _, err := newLeafValue(hashFunc, k, v)
1076 if err != nil {
1077 return false, err
1078 }
1079
1080 path := getPath(len(siblings), keyPath)
1081 for i := len(siblings) - 1; i >= 0; i-- {
1082 if path[i] {
1083 key, _, err = newIntermediate(hashFunc, siblings[i], key)
1084 if err != nil {
1085 return false, err
1086 }
1087 } else {
1088 key, _, err = newIntermediate(hashFunc, key, siblings[i])
1089 if err != nil {
1090 return false, err
1091 }
1092 }
1093 }
1094 if bytes.Equal(key[:], root) {
1095 return true, nil
1096 }
1097 return false, nil
1098}
1099
1100func (t *Tree) incNLeafs(wTx db.WriteTx, nLeafs int) error {
1101 oldNLeafs, err := t.GetNLeafsWithTx(wTx)
1102 if err != nil {
1103 return err
1104 }
1105 newNLeafs := oldNLeafs + nLeafs
1106 return t.setNLeafs(wTx, newNLeafs)
1107}
1108
1109func (t *Tree) setNLeafs(wTx db.WriteTx, nLeafs int) error {
1110 b := make([]byte, 8)
1111 binary.LittleEndian.PutUint64(b, uint64(nLeafs))
1112 if err := wTx.Set(dbKeyNLeafs, b); err != nil {
1113 return err
1114 }
1115 return nil
1116}
1117
1118// GetNLeafs returns the number of Leafs of the Tree.
1119func (t *Tree) GetNLeafs() (int, error) {
1120 rTx := t.db.ReadTx()
1121 defer rTx.Discard()
1122
1123 return t.GetNLeafsWithTx(rTx)
1124}
1125
1126// GetNLeafsWithTx does the same than the GetNLeafs method, but allowing to
1127// pass the db.ReadTx that is used.
1128func (t *Tree) GetNLeafsWithTx(rTx db.ReadTx) (int, error) {
1129 b, err := rTx.Get(dbKeyNLeafs)
1130 if err != nil {
1131 return 0, err
1132 }
1133 nLeafs := binary.LittleEndian.Uint64(b)
1134 return int(nLeafs), nil
1135}
1136
1137// SetRoot sets the root to the given root
1138func (t *Tree) SetRoot(root []byte) error {
1139 wTx := t.db.WriteTx()
1140 defer wTx.Discard()
1141
1142 if err := t.SetRootWithTx(wTx, root); err != nil {
1143 return err
1144 }
1145 return wTx.Commit()
1146}
1147
1148// SetRootWithTx sets the root to the given root using the given db.WriteTx
1149func (t *Tree) SetRootWithTx(wTx db.WriteTx, root []byte) error {
1150 if !t.editable() {
1151 return ErrSnapshotNotEditable
1152 }
1153
1154 if root == nil {
1155 return fmt.Errorf("can not SetRoot with nil root")
1156 }
1157
1158 // check that the root exists in the db
1159 if !bytes.Equal(root, t.emptyHash) {
1160 if _, err := wTx.Get(root); err == ErrKeyNotFound {
1161 return fmt.Errorf("can not SetRoot with root %x, as it"+
1162 " does not exist in the db", root)
1163 } else if err != nil {
1164 return err
1165 }
1166 }
1167
1168 return wTx.Set(dbKeyRoot, root)
1169}
1170
1171// Snapshot returns a read-only copy of the Tree from the given root
1172func (t *Tree) Snapshot(fromRoot []byte) (*Tree, error) {
1173 // allow to define which root to use
1174 if fromRoot == nil {
1175 var err error
1176 fromRoot, err = t.Root()
1177 if err != nil {
1178 return nil, err
1179 }
1180 }
1181 rTx := t.db.ReadTx()
1182 defer rTx.Discard()
1183 // check that the root exists in the db
1184 if !bytes.Equal(fromRoot, t.emptyHash) {
1185 if _, err := rTx.Get(fromRoot); err == ErrKeyNotFound {
1186 return nil,
1187 fmt.Errorf("can not do a Snapshot with root %x,"+
1188 " as it does not exist in the db", fromRoot)
1189 } else if err != nil {
1190 return nil, err
1191 }
1192 }
1193
1194 return &Tree{
1195 db: t.db,
1196 maxLevels: t.maxLevels,
1197 snapshotRoot: fromRoot,
1198 emptyHash: t.emptyHash,
1199 hashFunction: t.hashFunction,
1200 dbg: t.dbg,
1201 }, nil
1202}
1203
1204// Iterate iterates through the full Tree, executing the given function on each
1205// node of the Tree.
1206func (t *Tree) Iterate(fromRoot []byte, f func([]byte, []byte)) error {
1207 rTx := t.db.ReadTx()
1208 defer rTx.Discard()
1209
1210 return t.IterateWithTx(rTx, fromRoot, f)
1211}
1212
1213// IterateWithTx does the same than the Iterate method, but allowing to pass
1214// the db.ReadTx that is used.
1215func (t *Tree) IterateWithTx(rTx db.ReadTx, fromRoot []byte, f func([]byte, []byte)) error {
1216 // allow to define which root to use
1217 if fromRoot == nil {
1218 var err error
1219 fromRoot, err = t.RootWithTx(rTx)
1220 if err != nil {
1221 return err
1222 }
1223 }
1224 return t.iter(rTx, fromRoot, f)
1225}
1226
1227// IterateWithStop does the same than Iterate, but with int for the current
1228// level, and a boolean parameter used by the passed function, is to indicate to
1229// stop iterating on the branch when the method returns 'true'.
1230func (t *Tree) IterateWithStop(fromRoot []byte, f func(int, []byte, []byte) bool) error {
1231 rTx := t.db.ReadTx()
1232 defer rTx.Discard()
1233
1234 // allow to define which root to use
1235 if fromRoot == nil {
1236 var err error
1237 fromRoot, err = t.RootWithTx(rTx)
1238 if err != nil {
1239 return err
1240 }
1241 }
1242 return t.iterWithStop(rTx, fromRoot, 0, f)
1243}
1244
1245// IterateWithStopWithTx does the same than the IterateWithStop method, but
1246// allowing to pass the db.ReadTx that is used.
1247func (t *Tree) IterateWithStopWithTx(rTx db.ReadTx, fromRoot []byte,
1248 f func(int, []byte, []byte) bool) error {
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
1260func (t *Tree) iterWithStop(rTx db.ReadTx, k []byte, currLevel int,
1261 f func(int, []byte, []byte) bool) error {
1262 var v []byte
1263 var err error
1264 if bytes.Equal(k, t.emptyHash) {
1265 v = t.emptyHash
1266 } else {
1267 v, err = rTx.Get(k)
1268 if err != nil {
1269 return err
1270 }
1271 }
1272 currLevel++
1273
1274 switch v[0] {
1275 case PrefixValueEmpty:
1276 f(currLevel, k, v)
1277 case PrefixValueLeaf:
1278 f(currLevel, k, v)
1279 case PrefixValueIntermediate:
1280 stop := f(currLevel, k, v)
1281 if stop {
1282 return nil
1283 }
1284 l, r := ReadIntermediateChilds(v)
1285 if err = t.iterWithStop(rTx, l, currLevel, f); err != nil {
1286 return err
1287 }
1288 if err = t.iterWithStop(rTx, r, currLevel, f); err != nil {
1289 return err
1290 }
1291 default:
1292 return ErrInvalidValuePrefix
1293 }
1294 return nil
1295}
1296
1297func (t *Tree) iter(rTx db.ReadTx, k []byte, f func([]byte, []byte)) error {
1298 f2 := func(currLvl int, k, v []byte) bool {
1299 f(k, v)
1300 return false
1301 }
1302 return t.iterWithStop(rTx, k, 0, f2)
1303}
1304
1305// Dump exports all the Tree leafs in a byte array of length:
1306// [ N * (2+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
1307// [ 1 byte | 1 byte | S bytes | len(v) bytes ]
1308// [ len(k) | len(v) | key | value ]
1309// Where S is the size of the output of the hash function used for the Tree.
1310func (t *Tree) Dump(fromRoot []byte) ([]byte, error) {
1311 // allow to define which root to use
1312 if fromRoot == nil {
1313 var err error
1314 fromRoot, err = t.Root()
1315 if err != nil {
1316 return nil, err
1317 }
1318 }
1319
1320 // WARNING current encoding only supports key & values of 255 bytes each
1321 // (due using only 1 byte for the length headers).
1322 var b []byte
1323 var callbackErr error
1324 err := t.IterateWithStop(fromRoot, func(_ int, k, v []byte) bool {
1325 if v[0] != PrefixValueLeaf {
1326 return false
1327 }
1328 leafK, leafV := ReadLeafValue(v)
1329 kv := make([]byte, 2+len(leafK)+len(leafV))
1330 if len(leafK) > maxUint8 {
1331 callbackErr = fmt.Errorf("len(leafK) > %v", maxUint8)
1332 return true
1333 }
1334 kv[0] = byte(len(leafK))
1335 if len(leafV) > maxUint8 {
1336 callbackErr = fmt.Errorf("len(leafV) > %v", maxUint8)
1337 return true
1338 }
1339 kv[1] = byte(len(leafV))
1340 copy(kv[2:2+len(leafK)], leafK)
1341 copy(kv[2+len(leafK):], leafV)
1342 b = append(b, kv...)
1343 return false
1344 })
1345 if callbackErr != nil {
1346 return nil, callbackErr
1347 }
1348 return b, err
1349}
1350
1351// ImportDump imports the leafs (that have been exported with the Dump method)
1352// in the Tree.
1353func (t *Tree) ImportDump(b []byte) error {
1354 if !t.editable() {
1355 return ErrSnapshotNotEditable
1356 }
1357 root, err := t.Root()
1358 if err != nil {
1359 return err
1360 }
1361 if !bytes.Equal(root, t.emptyHash) {
1362 return ErrTreeNotEmpty
1363 }
1364
1365 r := bytes.NewReader(b)
1366 var keys, values [][]byte
1367 for {
1368 l := make([]byte, 2)
1369 _, err = io.ReadFull(r, l)
1370 if err == io.EOF {
1371 break
1372 } else if err != nil {
1373 return err
1374 }
1375 k := make([]byte, l[0])
1376 _, err = io.ReadFull(r, k)
1377 if err != nil {
1378 return err
1379 }
1380 v := make([]byte, l[1])
1381 _, err = io.ReadFull(r, v)
1382 if err != nil {
1383 return err
1384 }
1385 keys = append(keys, k)
1386 values = append(values, v)
1387 }
1388 if _, err = t.AddBatch(keys, values); err != nil {
1389 return err
1390 }
1391 return nil
1392}
1393
1394// Graphviz iterates across the full tree to generate a string Graphviz
1395// representation of the tree and writes it to w
1396func (t *Tree) Graphviz(w io.Writer, fromRoot []byte) error {
1397 return t.GraphvizFirstNLevels(w, fromRoot, t.maxLevels)
1398}
1399
1400// GraphvizFirstNLevels iterates across the first NLevels of the tree to
1401// generate a string Graphviz representation of the first NLevels of the tree
1402// and writes it to w
1403func (t *Tree) GraphvizFirstNLevels(w io.Writer, fromRoot []byte, untilLvl int) error {
1404 fmt.Fprintf(w, `digraph hierarchy {
1405node [fontname=Monospace,fontsize=10,shape=box]
1406`)
1407
1408 rTx := t.db.ReadTx()
1409 defer rTx.Discard()
1410
1411 if fromRoot == nil {
1412 var err error
1413 fromRoot, err = t.RootWithTx(rTx)
1414 if err != nil {
1415 return err
1416 }
1417 }
1418
1419 nEmpties := 0
1420 err := t.iterWithStop(rTx, fromRoot, 0, func(currLvl int, k, v []byte) bool {
1421 if currLvl == untilLvl {
1422 return true // to stop the iter from going down
1423 }
1424 switch v[0] {
1425 case PrefixValueEmpty:
1426 case PrefixValueLeaf:
1427 fmt.Fprintf(w, "\"%v\" [style=filled];\n", hex.EncodeToString(k[:nChars]))
1428 // key & value from the leaf
1429 kB, vB := ReadLeafValue(v)
1430 fmt.Fprintf(w, "\"%v\" -> {\"k:%v\\nv:%v\"}\n",
1431 hex.EncodeToString(k[:nChars]), hex.EncodeToString(kB[:nChars]),
1432 hex.EncodeToString(vB[:nChars]))
1433 fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
1434 hex.EncodeToString(kB[:nChars]), hex.EncodeToString(vB[:nChars]))
1435 case PrefixValueIntermediate:
1436 l, r := ReadIntermediateChilds(v)
1437 lStr := hex.EncodeToString(l[:nChars])
1438 rStr := hex.EncodeToString(r[:nChars])
1439 eStr := ""
1440 if bytes.Equal(l, t.emptyHash) {
1441 lStr = fmt.Sprintf("empty%v", nEmpties)
1442 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
1443 lStr)
1444 nEmpties++
1445 }
1446 if bytes.Equal(r, t.emptyHash) {
1447 rStr = fmt.Sprintf("empty%v", nEmpties)
1448 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
1449 rStr)
1450 nEmpties++
1451 }
1452 fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", hex.EncodeToString(k[:nChars]),
1453 lStr, rStr)
1454 fmt.Fprint(w, eStr)
1455 default:
1456 }
1457 return false
1458 })
1459 fmt.Fprintf(w, "}\n")
1460 return err
1461}
1462
1463// PrintGraphviz prints the output of Tree.Graphviz
1464func (t *Tree) PrintGraphviz(fromRoot []byte) error {
1465 if fromRoot == nil {
1466 var err error
1467 fromRoot, err = t.Root()
1468 if err != nil {
1469 return err
1470 }
1471 }
1472 return t.PrintGraphvizFirstNLevels(fromRoot, t.maxLevels)
1473}
1474
1475// PrintGraphvizFirstNLevels prints the output of Tree.GraphvizFirstNLevels
1476func (t *Tree) PrintGraphvizFirstNLevels(fromRoot []byte, untilLvl int) error {
1477 if fromRoot == nil {
1478 var err error
1479 fromRoot, err = t.Root()
1480 if err != nil {
1481 return err
1482 }
1483 }
1484 w := bytes.NewBufferString("")
1485 fmt.Fprintf(w,
1486 "--------\nGraphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+":\n")
1487 err := t.GraphvizFirstNLevels(w, fromRoot, untilLvl)
1488 if err != nil {
1489 fmt.Println(w)
1490 return err
1491 }
1492 fmt.Fprintf(w,
1493 "End of Graphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+"\n--------\n")
1494
1495 fmt.Println(w)
1496 return nil
1497}
1498
1499// TODO circom proofs
1500// TODO data structure for proofs (including root, key, value, siblings,
1501// hashFunction) + method to verify that data structure