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