arbo/tree.go
tree.goRaw
1/*
2Package arbo implements a Merkle Tree compatible with the circomlib
3implementation of the MerkleTree (when using the Poseidon hash function),
4following the specification from
5https://docs.iden3.io/publications/pdfs/Merkle-Tree.pdf and
6https://eprint.iacr.org/2018/955.
7
8Also allows to define which hash function to use. So for example, when working
9with zkSnarks the Poseidon hash function can be used, but when not, it can be
10used the Blake3 hash function, which improves the computation time.
11*/
12package arbo
13
14import (
15 "bytes"
16 "encoding/binary"
17 "encoding/hex"
18 "fmt"
19 "io"
20 "math"
21 "sync"
22
23 "github.com/iden3/go-merkletree/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
45var (
46 dbKeyRoot = []byte("root")
47 dbKeyNLeafs = []byte("nleafs")
48 emptyValue = []byte{0}
49
50 // ErrKeyAlreadyExists is used when trying to add a key as leaf to the
51 // tree that already exists.
52 ErrKeyAlreadyExists = fmt.Errorf("key already exists")
53 // ErrInvalidValuePrefix is used when going down into the tree, a value
54 // is read from the db and has an unrecognized prefix.
55 ErrInvalidValuePrefix = fmt.Errorf("invalid value prefix")
56 // ErrDBNoTx is used when trying to use Tree.dbPut but Tree.tx==nil
57 ErrDBNoTx = fmt.Errorf("dbPut error: no db Tx")
58 // ErrMaxLevel indicates when going down into the tree, the max level is
59 // reached
60 ErrMaxLevel = fmt.Errorf("max level reached")
61 // ErrMaxVirtualLevel indicates when going down into the tree, the max
62 // virtual level is reached
63 ErrMaxVirtualLevel = fmt.Errorf("max virtual level reached")
64)
65
66// Tree defines the struct that implements the MerkleTree functionalities
67type Tree struct {
68 sync.RWMutex
69 tx db.Tx
70 db db.Storage
71 maxLevels int
72 root []byte
73
74 hashFunction HashFunction
75 // TODO in the methods that use it, check if emptyHash param is len>0
76 // (check if it has been initialized)
77 emptyHash []byte
78
79 dbg *dbgStats
80}
81
82// NewTree returns a new Tree, if there is a Tree still in the given storage, it
83// will load it.
84func NewTree(storage db.Storage, maxLevels int, hash HashFunction) (*Tree, error) {
85 t := Tree{db: storage, maxLevels: maxLevels, hashFunction: hash}
86 t.emptyHash = make([]byte, t.hashFunction.Len()) // empty
87
88 root, err := t.dbGet(dbKeyRoot)
89 if err == db.ErrNotFound {
90 // store new root 0
91 t.tx, err = t.db.NewTx()
92 if err != nil {
93 return nil, err
94 }
95 t.root = t.emptyHash
96 if err = t.dbPut(dbKeyRoot, t.root); err != nil {
97 return nil, err
98 }
99 if err = t.setNLeafs(0); err != nil {
100 return nil, err
101 }
102 if err = t.tx.Commit(); err != nil {
103 return nil, err
104 }
105 return &t, err
106 } else if err != nil {
107 return nil, err
108 }
109 t.root = root
110 return &t, nil
111}
112
113// Root returns the root of the Tree
114func (t *Tree) Root() []byte {
115 return t.root
116}
117
118// HashFunction returns Tree.hashFunction
119func (t *Tree) HashFunction() HashFunction {
120 return t.hashFunction
121}
122
123// AddBatch adds a batch of key-values to the Tree. Returns an array containing
124// the indexes of the keys failed to add.
125func (t *Tree) AddBatch(keys, values [][]byte) ([]int, error) {
126 t.Lock()
127 defer t.Unlock()
128
129 vt, err := t.loadVT()
130 if err != nil {
131 return nil, err
132 }
133
134 // TODO check validity of keys & values for Tree.hashFunction
135
136 invalids, err := vt.addBatch(keys, values)
137 if err != nil {
138 return nil, err
139 }
140
141 // once the VirtualTree is build, compute the hashes
142 pairs, err := vt.computeHashes()
143 if err != nil {
144 // TODO currently invalids in computeHashes are not counted
145 return nil, err
146 }
147 t.root = vt.root.h
148
149 // store pairs in db
150 t.tx, err = t.db.NewTx()
151 if err != nil {
152 return nil, err
153 }
154 for i := 0; i < len(pairs); i++ {
155 if err := t.dbPut(pairs[i][0], pairs[i][1]); err != nil {
156 return nil, err
157 }
158 }
159
160 // store root to db
161 if err := t.dbPut(dbKeyRoot, t.root); err != nil {
162 return nil, err
163 }
164
165 // update nLeafs
166 if err := t.incNLeafs(len(keys) - len(invalids)); err != nil {
167 return nil, err
168 }
169
170 // commit db tx
171 if err := t.tx.Commit(); err != nil {
172 return nil, err
173 }
174 return invalids, nil
175}
176
177// loadVT loads a new virtual tree (vt) from the current Tree, which contains
178// the same leafs.
179func (t *Tree) loadVT() (vt, error) {
180 vt := newVT(t.maxLevels, t.hashFunction)
181 vt.params.dbg = t.dbg
182 err := t.Iterate(nil, func(k, v []byte) {
183 if v[0] != PrefixValueLeaf {
184 return
185 }
186 leafK, leafV := ReadLeafValue(v)
187 if err := vt.add(0, leafK, leafV); err != nil {
188 panic(err)
189 }
190 })
191
192 return vt, err
193}
194
195// Add inserts the key-value into the Tree. If the inputs come from a *big.Int,
196// is expected that are represented by a Little-Endian byte array (for circom
197// compatibility).
198func (t *Tree) Add(k, v []byte) error {
199 t.Lock()
200 defer t.Unlock()
201
202 var err error
203 t.tx, err = t.db.NewTx()
204 if err != nil {
205 return err
206 }
207
208 // TODO check validity of key & value for Tree.hashFunction
209
210 err = t.add(0, k, v) // add from level 0
211 if err != nil {
212 return err
213 }
214 // store root to db
215 if err := t.dbPut(dbKeyRoot, t.root); err != nil {
216 return err
217 }
218 // update nLeafs
219 if err = t.incNLeafs(1); err != nil {
220 return err
221 }
222 return t.tx.Commit()
223}
224
225func (t *Tree) add(fromLvl int, k, v []byte) error {
226 keyPath := make([]byte, t.hashFunction.Len())
227 copy(keyPath[:], k)
228
229 path := getPath(t.maxLevels, keyPath)
230 // go down to the leaf
231 var siblings [][]byte
232 _, _, siblings, err := t.down(k, t.root, siblings, path, fromLvl, false)
233 if err != nil {
234 return err
235 }
236
237 leafKey, leafValue, err := t.newLeafValue(k, v)
238 if err != nil {
239 return err
240 }
241
242 if err := t.dbPut(leafKey, leafValue); err != nil {
243 return err
244 }
245
246 // go up to the root
247 if len(siblings) == 0 {
248 t.root = leafKey
249 return nil
250 }
251 root, err := t.up(leafKey, siblings, path, len(siblings)-1, fromLvl)
252 if err != nil {
253 return err
254 }
255
256 t.root = root
257 return nil
258}
259
260// down goes down to the leaf recursively
261func (t *Tree) down(newKey, currKey []byte, siblings [][]byte,
262 path []bool, currLvl int, getLeaf bool) (
263 []byte, []byte, [][]byte, error) {
264 if currLvl > t.maxLevels-1 {
265 return nil, nil, nil, ErrMaxLevel
266 }
267
268 var err error
269 var currValue []byte
270 if bytes.Equal(currKey, t.emptyHash) {
271 // empty value
272 return currKey, emptyValue, siblings, nil
273 }
274 currValue, err = t.dbGet(currKey)
275 if err != nil {
276 return nil, nil, nil, err
277 }
278
279 switch currValue[0] {
280 case PrefixValueEmpty: // empty
281 // TODO WIP WARNING should not be reached, as the 'if' above should avoid
282 // reaching this point
283 // return currKey, empty, siblings, nil
284 panic("should not be reached, as the 'if' above should avoid reaching this point") // TMP
285 case PrefixValueLeaf: // leaf
286 if bytes.Equal(newKey, currKey) {
287 // TODO move this error msg to const & add test that
288 // checks that adding a repeated key this error is
289 // returned
290 return nil, nil, nil, ErrKeyAlreadyExists
291 }
292
293 if !bytes.Equal(currValue, emptyValue) {
294 if getLeaf {
295 return currKey, currValue, siblings, nil
296 }
297 oldLeafKey, _ := ReadLeafValue(currValue)
298 oldLeafKeyFull := make([]byte, t.hashFunction.Len())
299 copy(oldLeafKeyFull[:], oldLeafKey)
300
301 // if currKey is already used, go down until paths diverge
302 oldPath := getPath(t.maxLevels, oldLeafKeyFull)
303 siblings, err = t.downVirtually(siblings, currKey, newKey, oldPath, path, currLvl)
304 if err != nil {
305 return nil, nil, nil, err
306 }
307 }
308 return currKey, currValue, siblings, nil
309 case PrefixValueIntermediate: // intermediate
310 if len(currValue) != PrefixValueLen+t.hashFunction.Len()*2 {
311 return nil, nil, nil,
312 fmt.Errorf("intermediate value invalid length (expected: %d, actual: %d)",
313 PrefixValueLen+t.hashFunction.Len()*2, len(currValue))
314 }
315 // collect siblings while going down
316 if path[currLvl] {
317 // right
318 lChild, rChild := ReadIntermediateChilds(currValue)
319 siblings = append(siblings, lChild)
320 return t.down(newKey, rChild, siblings, path, currLvl+1, getLeaf)
321 }
322 // left
323 lChild, rChild := ReadIntermediateChilds(currValue)
324 siblings = append(siblings, rChild)
325 return t.down(newKey, lChild, siblings, path, currLvl+1, getLeaf)
326 default:
327 return nil, nil, nil, ErrInvalidValuePrefix
328 }
329}
330
331// downVirtually is used when in a leaf already exists, and a new leaf which
332// shares the path until the existing leaf is being added
333func (t *Tree) downVirtually(siblings [][]byte, oldKey, newKey []byte, oldPath,
334 newPath []bool, currLvl int) ([][]byte, error) {
335 var err error
336 if currLvl > t.maxLevels-1 {
337 return nil, ErrMaxVirtualLevel
338 }
339
340 if oldPath[currLvl] == newPath[currLvl] {
341 siblings = append(siblings, t.emptyHash)
342
343 siblings, err = t.downVirtually(siblings, oldKey, newKey, oldPath, newPath, currLvl+1)
344 if err != nil {
345 return nil, err
346 }
347 return siblings, nil
348 }
349 // reached the divergence
350 siblings = append(siblings, oldKey)
351
352 return siblings, nil
353}
354
355// up goes up recursively updating the intermediate nodes
356func (t *Tree) up(key []byte, siblings [][]byte, path []bool, currLvl, toLvl int) ([]byte, error) {
357 var k, v []byte
358 var err error
359 if path[currLvl+toLvl] {
360 k, v, err = t.newIntermediate(siblings[currLvl], key)
361 if err != nil {
362 return nil, err
363 }
364 } else {
365 k, v, err = t.newIntermediate(key, siblings[currLvl])
366 if err != nil {
367 return nil, err
368 }
369 }
370 // store k-v to db
371 if err = t.dbPut(k, v); err != nil {
372 return nil, err
373 }
374
375 if currLvl == 0 {
376 // reached the root
377 return k, nil
378 }
379
380 return t.up(k, siblings, path, currLvl-1, toLvl)
381}
382
383func (t *Tree) newLeafValue(k, v []byte) ([]byte, []byte, error) {
384 t.dbg.incHash()
385 return newLeafValue(t.hashFunction, k, v)
386}
387
388func newLeafValue(hashFunc HashFunction, k, v []byte) ([]byte, []byte, error) {
389 leafKey, err := hashFunc.Hash(k, v, []byte{1})
390 if err != nil {
391 return nil, nil, err
392 }
393 var leafValue []byte
394 leafValue = append(leafValue, byte(1))
395 leafValue = append(leafValue, byte(len(k)))
396 leafValue = append(leafValue, k...)
397 leafValue = append(leafValue, v...)
398 return leafKey, leafValue, nil
399}
400
401// ReadLeafValue reads from a byte array the leaf key & value
402func ReadLeafValue(b []byte) ([]byte, []byte) {
403 if len(b) < PrefixValueLen {
404 return []byte{}, []byte{}
405 }
406
407 kLen := b[1]
408 if len(b) < PrefixValueLen+int(kLen) {
409 return []byte{}, []byte{}
410 }
411 k := b[PrefixValueLen : PrefixValueLen+kLen]
412 v := b[PrefixValueLen+kLen:]
413 return k, v
414}
415
416func (t *Tree) newIntermediate(l, r []byte) ([]byte, []byte, error) {
417 t.dbg.incHash()
418 return newIntermediate(t.hashFunction, l, r)
419}
420
421func newIntermediate(hashFunc HashFunction, l, r []byte) ([]byte, []byte, error) {
422 b := make([]byte, PrefixValueLen+hashFunc.Len()*2)
423 b[0] = 2
424 b[1] = byte(len(l))
425 copy(b[PrefixValueLen:PrefixValueLen+hashFunc.Len()], l)
426 copy(b[PrefixValueLen+hashFunc.Len():], r)
427
428 key, err := hashFunc.Hash(l, r)
429 if err != nil {
430 return nil, nil, err
431 }
432
433 return key, b, nil
434}
435
436// ReadIntermediateChilds reads from a byte array the two childs keys
437func ReadIntermediateChilds(b []byte) ([]byte, []byte) {
438 if len(b) < PrefixValueLen {
439 return []byte{}, []byte{}
440 }
441
442 lLen := b[1]
443 if len(b) < PrefixValueLen+int(lLen) {
444 return []byte{}, []byte{}
445 }
446 l := b[PrefixValueLen : PrefixValueLen+lLen]
447 r := b[PrefixValueLen+lLen:]
448 return l, r
449}
450
451func getPath(numLevels int, k []byte) []bool {
452 path := make([]bool, numLevels)
453 for n := 0; n < numLevels; n++ {
454 path[n] = k[n/8]&(1<<(n%8)) != 0
455 }
456 return path
457}
458
459// Update updates the value for a given existing key. If the given key does not
460// exist, returns an error.
461func (t *Tree) Update(k, v []byte) error {
462 t.Lock()
463 defer t.Unlock()
464
465 var err error
466 t.tx, err = t.db.NewTx()
467 if err != nil {
468 return err
469 }
470
471 keyPath := make([]byte, t.hashFunction.Len())
472 copy(keyPath[:], k)
473 path := getPath(t.maxLevels, keyPath)
474
475 var siblings [][]byte
476 _, valueAtBottom, siblings, err := t.down(k, t.root, siblings, path, 0, true)
477 if err != nil {
478 return err
479 }
480 oldKey, _ := ReadLeafValue(valueAtBottom)
481 if !bytes.Equal(oldKey, k) {
482 return fmt.Errorf("key %s does not exist", hex.EncodeToString(k))
483 }
484
485 leafKey, leafValue, err := t.newLeafValue(k, v)
486 if err != nil {
487 return err
488 }
489
490 if err := t.dbPut(leafKey, leafValue); err != nil {
491 return err
492 }
493
494 // go up to the root
495 if len(siblings) == 0 {
496 t.root = leafKey
497 return t.tx.Commit()
498 }
499 root, err := t.up(leafKey, siblings, path, len(siblings)-1, 0)
500 if err != nil {
501 return err
502 }
503
504 t.root = root
505 // store root to db
506 if err := t.dbPut(dbKeyRoot, t.root); err != nil {
507 return err
508 }
509 return t.tx.Commit()
510}
511
512// GenProof generates a MerkleTree proof for the given key. If the key exists in
513// the Tree, the proof will be of existence, if the key does not exist in the
514// tree, the proof will be of non-existence.
515func (t *Tree) GenProof(k []byte) ([]byte, []byte, error) {
516 keyPath := make([]byte, t.hashFunction.Len())
517 copy(keyPath[:], k)
518
519 path := getPath(t.maxLevels, keyPath)
520 // go down to the leaf
521 var siblings [][]byte
522 _, value, siblings, err := t.down(k, t.root, siblings, path, 0, true)
523 if err != nil {
524 return nil, nil, err
525 }
526
527 leafK, leafV := ReadLeafValue(value)
528 if !bytes.Equal(k, leafK) {
529 fmt.Println("key not in Tree")
530 fmt.Println(leafK)
531 fmt.Println(leafV)
532 // TODO proof of non-existence
533 panic("unimplemented")
534 }
535
536 s := PackSiblings(t.hashFunction, siblings)
537 return leafV, s, nil
538}
539
540// PackSiblings packs the siblings into a byte array.
541// [ 1 byte | L bytes | S * N bytes ]
542// [ bitmap length (L) | bitmap | N non-zero siblings ]
543// Where the bitmap indicates if the sibling is 0 or a value from the siblings
544// array. And S is the size of the output of the hash function used for the
545// Tree.
546func PackSiblings(hashFunc HashFunction, siblings [][]byte) []byte {
547 var b []byte
548 var bitmap []bool
549 emptySibling := make([]byte, hashFunc.Len())
550 for i := 0; i < len(siblings); i++ {
551 if bytes.Equal(siblings[i], emptySibling) {
552 bitmap = append(bitmap, false)
553 } else {
554 bitmap = append(bitmap, true)
555 b = append(b, siblings[i]...)
556 }
557 }
558
559 bitmapBytes := bitmapToBytes(bitmap)
560 l := len(bitmapBytes)
561
562 res := make([]byte, l+1+len(b))
563 res[0] = byte(l) // set the bitmapBytes length
564 copy(res[1:1+l], bitmapBytes)
565 copy(res[1+l:], b)
566 return res
567}
568
569// UnpackSiblings unpacks the siblings from a byte array.
570func UnpackSiblings(hashFunc HashFunction, b []byte) ([][]byte, error) {
571 l := b[0]
572 bitmapBytes := b[1 : 1+l]
573 bitmap := bytesToBitmap(bitmapBytes)
574 siblingsBytes := b[1+l:]
575 iSibl := 0
576 emptySibl := make([]byte, hashFunc.Len())
577 var siblings [][]byte
578 for i := 0; i < len(bitmap); i++ {
579 if iSibl >= len(siblingsBytes) {
580 break
581 }
582 if bitmap[i] {
583 siblings = append(siblings, siblingsBytes[iSibl:iSibl+hashFunc.Len()])
584 iSibl += hashFunc.Len()
585 } else {
586 siblings = append(siblings, emptySibl)
587 }
588 }
589 return siblings, nil
590}
591
592func bitmapToBytes(bitmap []bool) []byte {
593 bitmapBytesLen := int(math.Ceil(float64(len(bitmap)) / 8)) //nolint:gomnd
594 b := make([]byte, bitmapBytesLen)
595 for i := 0; i < len(bitmap); i++ {
596 if bitmap[i] {
597 b[i/8] |= 1 << (i % 8)
598 }
599 }
600 return b
601}
602
603func bytesToBitmap(b []byte) []bool {
604 var bitmap []bool
605 for i := 0; i < len(b); i++ {
606 for j := 0; j < 8; j++ {
607 bitmap = append(bitmap, b[i]&(1<<j) > 0)
608 }
609 }
610 return bitmap
611}
612
613// Get returns the value for a given key
614func (t *Tree) Get(k []byte) ([]byte, []byte, error) {
615 keyPath := make([]byte, t.hashFunction.Len())
616 copy(keyPath[:], k)
617
618 path := getPath(t.maxLevels, keyPath)
619 // go down to the leaf
620 var siblings [][]byte
621 _, value, _, err := t.down(k, t.root, siblings, path, 0, true)
622 if err != nil {
623 return nil, nil, err
624 }
625 leafK, leafV := ReadLeafValue(value)
626 if !bytes.Equal(k, leafK) {
627 return leafK, leafV, fmt.Errorf("Tree.Get error: keys doesn't match, %s != %s",
628 BytesToBigInt(k), BytesToBigInt(leafK))
629 }
630
631 return leafK, leafV, nil
632}
633
634// CheckProof verifies the given proof. The proof verification depends on the
635// HashFunction passed as parameter.
636func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool, error) {
637 siblings, err := UnpackSiblings(hashFunc, packedSiblings)
638 if err != nil {
639 return false, err
640 }
641
642 keyPath := make([]byte, hashFunc.Len())
643 copy(keyPath[:], k)
644
645 key, _, err := newLeafValue(hashFunc, k, v)
646 if err != nil {
647 return false, err
648 }
649
650 path := getPath(len(siblings), keyPath)
651 for i := len(siblings) - 1; i >= 0; i-- {
652 if path[i] {
653 key, _, err = newIntermediate(hashFunc, siblings[i], key)
654 if err != nil {
655 return false, err
656 }
657 } else {
658 key, _, err = newIntermediate(hashFunc, key, siblings[i])
659 if err != nil {
660 return false, err
661 }
662 }
663 }
664 if bytes.Equal(key[:], root) {
665 return true, nil
666 }
667 return false, nil
668}
669
670func (t *Tree) dbPut(k, v []byte) error {
671 if t.tx == nil {
672 return ErrDBNoTx
673 }
674 t.dbg.incDbPut()
675 return t.tx.Put(k, v)
676}
677
678func (t *Tree) dbGet(k []byte) ([]byte, error) {
679 // if key is empty, return empty as value
680 if bytes.Equal(k, t.emptyHash) {
681 return t.emptyHash, nil
682 }
683 t.dbg.incDbGet()
684
685 v, err := t.db.Get(k)
686 if err == nil {
687 return v, nil
688 }
689 if t.tx != nil {
690 return t.tx.Get(k)
691 }
692 return nil, db.ErrNotFound
693}
694
695// Warning: should be called with a Tree.tx created, and with a Tree.tx.Commit
696// after the setNLeafs call.
697func (t *Tree) incNLeafs(nLeafs int) error {
698 oldNLeafs, err := t.GetNLeafs()
699 if err != nil {
700 return err
701 }
702 newNLeafs := oldNLeafs + nLeafs
703 return t.setNLeafs(newNLeafs)
704}
705
706// Warning: should be called with a Tree.tx created, and with a Tree.tx.Commit
707// after the setNLeafs call.
708func (t *Tree) setNLeafs(nLeafs int) error {
709 b := make([]byte, 8)
710 binary.LittleEndian.PutUint64(b, uint64(nLeafs))
711 if err := t.dbPut(dbKeyNLeafs, b); err != nil {
712 return err
713 }
714 return nil
715}
716
717// GetNLeafs returns the number of Leafs of the Tree.
718func (t *Tree) GetNLeafs() (int, error) {
719 b, err := t.dbGet(dbKeyNLeafs)
720 if err != nil {
721 return 0, err
722 }
723 nLeafs := binary.LittleEndian.Uint64(b)
724 return int(nLeafs), nil
725}
726
727// Iterate iterates through the full Tree, executing the given function on each
728// node of the Tree.
729func (t *Tree) Iterate(rootKey []byte, f func([]byte, []byte)) error {
730 // allow to define which root to use
731 if rootKey == nil {
732 rootKey = t.Root()
733 }
734 return t.iter(rootKey, f)
735}
736
737// IterateWithStop does the same than Iterate, but with int for the current
738// level, and a boolean parameter used by the passed function, is to indicate to
739// stop iterating on the branch when the method returns 'true'.
740func (t *Tree) IterateWithStop(rootKey []byte, f func(int, []byte, []byte) bool) error {
741 // allow to define which root to use
742 if rootKey == nil {
743 rootKey = t.Root()
744 }
745 return t.iterWithStop(rootKey, 0, f)
746}
747
748func (t *Tree) iterWithStop(k []byte, currLevel int, f func(int, []byte, []byte) bool) error {
749 v, err := t.dbGet(k)
750 if err != nil {
751 return err
752 }
753 currLevel++
754
755 switch v[0] {
756 case PrefixValueEmpty:
757 f(currLevel, k, v)
758 case PrefixValueLeaf:
759 f(currLevel, k, v)
760 case PrefixValueIntermediate:
761 stop := f(currLevel, k, v)
762 if stop {
763 return nil
764 }
765 l, r := ReadIntermediateChilds(v)
766 if err = t.iterWithStop(l, currLevel, f); err != nil {
767 return err
768 }
769 if err = t.iterWithStop(r, currLevel, f); err != nil {
770 return err
771 }
772 default:
773 return ErrInvalidValuePrefix
774 }
775 return nil
776}
777
778func (t *Tree) iter(k []byte, f func([]byte, []byte)) error {
779 f2 := func(currLvl int, k, v []byte) bool {
780 f(k, v)
781 return false
782 }
783 return t.iterWithStop(k, 0, f2)
784}
785
786// Dump exports all the Tree leafs in a byte array of length:
787// [ N * (2+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
788// [ 1 byte | 1 byte | S bytes | len(v) bytes ]
789// [ len(k) | len(v) | key | value ]
790// Where S is the size of the output of the hash function used for the Tree.
791func (t *Tree) Dump(rootKey []byte) ([]byte, error) {
792 // allow to define which root to use
793 if rootKey == nil {
794 rootKey = t.Root()
795 }
796
797 // WARNING current encoding only supports key & values of 255 bytes each
798 // (due using only 1 byte for the length headers).
799 var b []byte
800 err := t.Iterate(rootKey, func(k, v []byte) {
801 if v[0] != PrefixValueLeaf {
802 return
803 }
804 leafK, leafV := ReadLeafValue(v)
805 kv := make([]byte, 2+len(leafK)+len(leafV))
806 kv[0] = byte(len(leafK))
807 kv[1] = byte(len(leafV))
808 copy(kv[2:2+len(leafK)], leafK)
809 copy(kv[2+len(leafK):], leafV)
810 b = append(b, kv...)
811 })
812 return b, err
813}
814
815// ImportDump imports the leafs (that have been exported with the ExportLeafs
816// method) in the Tree.
817func (t *Tree) ImportDump(b []byte) error {
818 r := bytes.NewReader(b)
819 var err error
820 var keys, values [][]byte
821 for {
822 l := make([]byte, 2)
823 _, err = io.ReadFull(r, l)
824 if err == io.EOF {
825 break
826 } else if err != nil {
827 return err
828 }
829 k := make([]byte, l[0])
830 _, err = io.ReadFull(r, k)
831 if err != nil {
832 return err
833 }
834 v := make([]byte, l[1])
835 _, err = io.ReadFull(r, v)
836 if err != nil {
837 return err
838 }
839 keys = append(keys, k)
840 values = append(values, v)
841 }
842 if _, err = t.AddBatch(keys, values); err != nil {
843 return err
844 }
845 return nil
846}
847
848// Graphviz iterates across the full tree to generate a string Graphviz
849// representation of the tree and writes it to w
850func (t *Tree) Graphviz(w io.Writer, rootKey []byte) error {
851 return t.GraphvizFirstNLevels(w, rootKey, t.maxLevels)
852}
853
854// GraphvizFirstNLevels iterates across the first NLevels of the tree to
855// generate a string Graphviz representation of the first NLevels of the tree
856// and writes it to w
857func (t *Tree) GraphvizFirstNLevels(w io.Writer, rootKey []byte, untilLvl int) error {
858 fmt.Fprintf(w, `digraph hierarchy {
859node [fontname=Monospace,fontsize=10,shape=box]
860`)
861 if rootKey == nil {
862 rootKey = t.Root()
863 }
864 nEmpties := 0
865 err := t.iterWithStop(rootKey, 0, func(currLvl int, k, v []byte) bool {
866 if currLvl == untilLvl {
867 return true // to stop the iter from going down
868 }
869 switch v[0] {
870 case PrefixValueEmpty:
871 case PrefixValueLeaf:
872 fmt.Fprintf(w, "\"%v\" [style=filled];\n", hex.EncodeToString(k[:nChars]))
873 // key & value from the leaf
874 kB, vB := ReadLeafValue(v)
875 fmt.Fprintf(w, "\"%v\" -> {\"k:%v\\nv:%v\"}\n",
876 hex.EncodeToString(k[:nChars]), hex.EncodeToString(kB[:nChars]),
877 hex.EncodeToString(vB[:nChars]))
878 fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
879 hex.EncodeToString(kB[:nChars]), hex.EncodeToString(vB[:nChars]))
880 case PrefixValueIntermediate:
881 l, r := ReadIntermediateChilds(v)
882 lStr := hex.EncodeToString(l[:nChars])
883 rStr := hex.EncodeToString(r[:nChars])
884 eStr := ""
885 if bytes.Equal(l, t.emptyHash) {
886 lStr = fmt.Sprintf("empty%v", nEmpties)
887 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
888 lStr)
889 nEmpties++
890 }
891 if bytes.Equal(r, t.emptyHash) {
892 rStr = fmt.Sprintf("empty%v", nEmpties)
893 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
894 rStr)
895 nEmpties++
896 }
897 fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", hex.EncodeToString(k[:nChars]),
898 lStr, rStr)
899 fmt.Fprint(w, eStr)
900 default:
901 }
902 return false
903 })
904 fmt.Fprintf(w, "}\n")
905 return err
906}
907
908// PrintGraphviz prints the output of Tree.Graphviz
909func (t *Tree) PrintGraphviz(rootKey []byte) error {
910 if rootKey == nil {
911 rootKey = t.Root()
912 }
913 return t.PrintGraphvizFirstNLevels(rootKey, t.maxLevels)
914}
915
916// PrintGraphvizFirstNLevels prints the output of Tree.GraphvizFirstNLevels
917func (t *Tree) PrintGraphvizFirstNLevels(rootKey []byte, untilLvl int) error {
918 if rootKey == nil {
919 rootKey = t.Root()
920 }
921 w := bytes.NewBufferString("")
922 fmt.Fprintf(w,
923 "--------\nGraphviz of the Tree with Root "+hex.EncodeToString(rootKey)+":\n")
924 err := t.GraphvizFirstNLevels(w, rootKey, untilLvl)
925 if err != nil {
926 fmt.Println(w)
927 return err
928 }
929 fmt.Fprintf(w,
930 "End of Graphviz of the Tree with Root "+hex.EncodeToString(rootKey)+"\n--------\n")
931
932 fmt.Println(w)
933 return nil
934}
935
936// Purge WIP: unimplemented TODO
937func (t *Tree) Purge(keys [][]byte) error {
938 return nil
939}
940
941// TODO circom proofs