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 fmt.Printf("newKey: %s, currKey: %s, currLvl: %d, currValue: %s\n",
282 hex.EncodeToString(newKey), hex.EncodeToString(currKey),
283 currLvl, hex.EncodeToString(currValue))
284 panic("This point should not be reached, as the 'if' above" +
285 " should avoid reaching this point. This panic is temporary" +
286 " for reporting purposes, will be deleted in future versions." +
287 " Please paste this log (including the previous lines) in a" +
288 " new issue: https://github.com/arnaucube/arbo/issues/new") // TMP
289 case PrefixValueLeaf: // leaf
290 if !bytes.Equal(currValue, emptyValue) {
291 if getLeaf {
292 return currKey, currValue, siblings, nil
293 }
294 oldLeafKey, _ := ReadLeafValue(currValue)
295 if bytes.Equal(newKey, oldLeafKey) {
296 return nil, nil, nil, ErrKeyAlreadyExists
297 }
298
299 oldLeafKeyFull := make([]byte, t.hashFunction.Len())
300 copy(oldLeafKeyFull[:], oldLeafKey)
301
302 // if currKey is already used, go down until paths diverge
303 oldPath := getPath(t.maxLevels, oldLeafKeyFull)
304 siblings, err = t.downVirtually(siblings, currKey, newKey, oldPath, path, currLvl)
305 if err != nil {
306 return nil, nil, nil, err
307 }
308 }
309 return currKey, currValue, siblings, nil
310 case PrefixValueIntermediate: // intermediate
311 if len(currValue) != PrefixValueLen+t.hashFunction.Len()*2 {
312 return nil, nil, nil,
313 fmt.Errorf("intermediate value invalid length (expected: %d, actual: %d)",
314 PrefixValueLen+t.hashFunction.Len()*2, len(currValue))
315 }
316 // collect siblings while going down
317 if path[currLvl] {
318 // right
319 lChild, rChild := ReadIntermediateChilds(currValue)
320 siblings = append(siblings, lChild)
321 return t.down(newKey, rChild, siblings, path, currLvl+1, getLeaf)
322 }
323 // left
324 lChild, rChild := ReadIntermediateChilds(currValue)
325 siblings = append(siblings, rChild)
326 return t.down(newKey, lChild, siblings, path, currLvl+1, getLeaf)
327 default:
328 return nil, nil, nil, ErrInvalidValuePrefix
329 }
330}
331
332// downVirtually is used when in a leaf already exists, and a new leaf which
333// shares the path until the existing leaf is being added
334func (t *Tree) downVirtually(siblings [][]byte, oldKey, newKey []byte, oldPath,
335 newPath []bool, currLvl int) ([][]byte, error) {
336 var err error
337 if currLvl > t.maxLevels-1 {
338 return nil, ErrMaxVirtualLevel
339 }
340
341 if oldPath[currLvl] == newPath[currLvl] {
342 siblings = append(siblings, t.emptyHash)
343
344 siblings, err = t.downVirtually(siblings, oldKey, newKey, oldPath, newPath, currLvl+1)
345 if err != nil {
346 return nil, err
347 }
348 return siblings, nil
349 }
350 // reached the divergence
351 siblings = append(siblings, oldKey)
352
353 return siblings, nil
354}
355
356// up goes up recursively updating the intermediate nodes
357func (t *Tree) up(key []byte, siblings [][]byte, path []bool, currLvl, toLvl int) ([]byte, error) {
358 var k, v []byte
359 var err error
360 if path[currLvl+toLvl] {
361 k, v, err = t.newIntermediate(siblings[currLvl], key)
362 if err != nil {
363 return nil, err
364 }
365 } else {
366 k, v, err = t.newIntermediate(key, siblings[currLvl])
367 if err != nil {
368 return nil, err
369 }
370 }
371 // store k-v to db
372 if err = t.dbPut(k, v); err != nil {
373 return nil, err
374 }
375
376 if currLvl == 0 {
377 // reached the root
378 return k, nil
379 }
380
381 return t.up(k, siblings, path, currLvl-1, toLvl)
382}
383
384func (t *Tree) newLeafValue(k, v []byte) ([]byte, []byte, error) {
385 t.dbg.incHash()
386 return newLeafValue(t.hashFunction, k, v)
387}
388
389// newLeafValue takes a key & value from a leaf, and computes the leaf hash,
390// which is used as the leaf key. And the value is the concatenation of the
391// inputed key & value. The output of this function is used as key-value to
392// store the leaf in the DB.
393// [ 1 byte | 1 byte | N bytes | M bytes ]
394// [ type of node | length of key | key | value ]
395func newLeafValue(hashFunc HashFunction, k, v []byte) ([]byte, []byte, error) {
396 leafKey, err := hashFunc.Hash(k, v, []byte{1})
397 if err != nil {
398 return nil, nil, err
399 }
400 var leafValue []byte
401 leafValue = append(leafValue, byte(1))
402 leafValue = append(leafValue, byte(len(k)))
403 leafValue = append(leafValue, k...)
404 leafValue = append(leafValue, v...)
405 return leafKey, leafValue, nil
406}
407
408// ReadLeafValue reads from a byte array the leaf key & value
409func ReadLeafValue(b []byte) ([]byte, []byte) {
410 if len(b) < PrefixValueLen {
411 return []byte{}, []byte{}
412 }
413
414 kLen := b[1]
415 if len(b) < PrefixValueLen+int(kLen) {
416 return []byte{}, []byte{}
417 }
418 k := b[PrefixValueLen : PrefixValueLen+kLen]
419 v := b[PrefixValueLen+kLen:]
420 return k, v
421}
422
423func (t *Tree) newIntermediate(l, r []byte) ([]byte, []byte, error) {
424 t.dbg.incHash()
425 return newIntermediate(t.hashFunction, l, r)
426}
427
428// newIntermediate takes the left & right keys of a intermediate node, and
429// computes its hash. Returns the hash of the node, which is the node key, and a
430// byte array that contains the value (which contains the left & right child
431// keys) to store in the DB.
432// [ 1 byte | 1 byte | N bytes | N bytes ]
433// [ type of node | length of key | left key | right key ]
434func newIntermediate(hashFunc HashFunction, l, r []byte) ([]byte, []byte, error) {
435 b := make([]byte, PrefixValueLen+hashFunc.Len()*2)
436 b[0] = 2
437 b[1] = byte(len(l))
438 copy(b[PrefixValueLen:PrefixValueLen+hashFunc.Len()], l)
439 copy(b[PrefixValueLen+hashFunc.Len():], r)
440
441 key, err := hashFunc.Hash(l, r)
442 if err != nil {
443 return nil, nil, err
444 }
445
446 return key, b, nil
447}
448
449// ReadIntermediateChilds reads from a byte array the two childs keys
450func ReadIntermediateChilds(b []byte) ([]byte, []byte) {
451 if len(b) < PrefixValueLen {
452 return []byte{}, []byte{}
453 }
454
455 lLen := b[1]
456 if len(b) < PrefixValueLen+int(lLen) {
457 return []byte{}, []byte{}
458 }
459 l := b[PrefixValueLen : PrefixValueLen+lLen]
460 r := b[PrefixValueLen+lLen:]
461 return l, r
462}
463
464func getPath(numLevels int, k []byte) []bool {
465 path := make([]bool, numLevels)
466 for n := 0; n < numLevels; n++ {
467 path[n] = k[n/8]&(1<<(n%8)) != 0
468 }
469 return path
470}
471
472// Update updates the value for a given existing key. If the given key does not
473// exist, returns an error.
474func (t *Tree) Update(k, v []byte) error {
475 t.Lock()
476 defer t.Unlock()
477
478 var err error
479 t.tx, err = t.db.NewTx()
480 if err != nil {
481 return err
482 }
483
484 keyPath := make([]byte, t.hashFunction.Len())
485 copy(keyPath[:], k)
486 path := getPath(t.maxLevels, keyPath)
487
488 var siblings [][]byte
489 _, valueAtBottom, siblings, err := t.down(k, t.root, siblings, path, 0, true)
490 if err != nil {
491 return err
492 }
493 oldKey, _ := ReadLeafValue(valueAtBottom)
494 if !bytes.Equal(oldKey, k) {
495 return fmt.Errorf("key %s does not exist", hex.EncodeToString(k))
496 }
497
498 leafKey, leafValue, err := t.newLeafValue(k, v)
499 if err != nil {
500 return err
501 }
502
503 if err := t.dbPut(leafKey, leafValue); err != nil {
504 return err
505 }
506
507 // go up to the root
508 if len(siblings) == 0 {
509 t.root = leafKey
510 return t.tx.Commit()
511 }
512 root, err := t.up(leafKey, siblings, path, len(siblings)-1, 0)
513 if err != nil {
514 return err
515 }
516
517 t.root = root
518 // store root to db
519 if err := t.dbPut(dbKeyRoot, t.root); err != nil {
520 return err
521 }
522 return t.tx.Commit()
523}
524
525// GenProof generates a MerkleTree proof for the given key. If the key exists in
526// the Tree, the proof will be of existence, if the key does not exist in the
527// tree, the proof will be of non-existence.
528func (t *Tree) GenProof(k []byte) ([]byte, []byte, error) {
529 keyPath := make([]byte, t.hashFunction.Len())
530 copy(keyPath[:], k)
531
532 path := getPath(t.maxLevels, keyPath)
533 // go down to the leaf
534 var siblings [][]byte
535 _, value, siblings, err := t.down(k, t.root, siblings, path, 0, true)
536 if err != nil {
537 return nil, nil, err
538 }
539
540 leafK, leafV := ReadLeafValue(value)
541 if !bytes.Equal(k, leafK) {
542 fmt.Println("key not in Tree")
543 fmt.Println(leafK)
544 fmt.Println(leafV)
545 // TODO proof of non-existence
546 panic("unimplemented")
547 }
548
549 s := PackSiblings(t.hashFunction, siblings)
550 return leafV, s, nil
551}
552
553// PackSiblings packs the siblings into a byte array.
554// [ 1 byte | L bytes | S * N bytes ]
555// [ bitmap length (L) | bitmap | N non-zero siblings ]
556// Where the bitmap indicates if the sibling is 0 or a value from the siblings
557// array. And S is the size of the output of the hash function used for the
558// Tree.
559func PackSiblings(hashFunc HashFunction, siblings [][]byte) []byte {
560 var b []byte
561 var bitmap []bool
562 emptySibling := make([]byte, hashFunc.Len())
563 for i := 0; i < len(siblings); i++ {
564 if bytes.Equal(siblings[i], emptySibling) {
565 bitmap = append(bitmap, false)
566 } else {
567 bitmap = append(bitmap, true)
568 b = append(b, siblings[i]...)
569 }
570 }
571
572 bitmapBytes := bitmapToBytes(bitmap)
573 l := len(bitmapBytes)
574
575 res := make([]byte, l+1+len(b))
576 res[0] = byte(l) // set the bitmapBytes length
577 copy(res[1:1+l], bitmapBytes)
578 copy(res[1+l:], b)
579 return res
580}
581
582// UnpackSiblings unpacks the siblings from a byte array.
583func UnpackSiblings(hashFunc HashFunction, b []byte) ([][]byte, error) {
584 l := b[0]
585 bitmapBytes := b[1 : 1+l]
586 bitmap := bytesToBitmap(bitmapBytes)
587 siblingsBytes := b[1+l:]
588 iSibl := 0
589 emptySibl := make([]byte, hashFunc.Len())
590 var siblings [][]byte
591 for i := 0; i < len(bitmap); i++ {
592 if iSibl >= len(siblingsBytes) {
593 break
594 }
595 if bitmap[i] {
596 siblings = append(siblings, siblingsBytes[iSibl:iSibl+hashFunc.Len()])
597 iSibl += hashFunc.Len()
598 } else {
599 siblings = append(siblings, emptySibl)
600 }
601 }
602 return siblings, nil
603}
604
605func bitmapToBytes(bitmap []bool) []byte {
606 bitmapBytesLen := int(math.Ceil(float64(len(bitmap)) / 8)) //nolint:gomnd
607 b := make([]byte, bitmapBytesLen)
608 for i := 0; i < len(bitmap); i++ {
609 if bitmap[i] {
610 b[i/8] |= 1 << (i % 8)
611 }
612 }
613 return b
614}
615
616func bytesToBitmap(b []byte) []bool {
617 var bitmap []bool
618 for i := 0; i < len(b); i++ {
619 for j := 0; j < 8; j++ {
620 bitmap = append(bitmap, b[i]&(1<<j) > 0)
621 }
622 }
623 return bitmap
624}
625
626// Get returns the value for a given key
627func (t *Tree) Get(k []byte) ([]byte, []byte, error) {
628 keyPath := make([]byte, t.hashFunction.Len())
629 copy(keyPath[:], k)
630
631 path := getPath(t.maxLevels, keyPath)
632 // go down to the leaf
633 var siblings [][]byte
634 _, value, _, err := t.down(k, t.root, siblings, path, 0, true)
635 if err != nil {
636 return nil, nil, err
637 }
638 leafK, leafV := ReadLeafValue(value)
639 if !bytes.Equal(k, leafK) {
640 return leafK, leafV, fmt.Errorf("Tree.Get error: keys doesn't match, %s != %s",
641 BytesToBigInt(k), BytesToBigInt(leafK))
642 }
643
644 return leafK, leafV, nil
645}
646
647// CheckProof verifies the given proof. The proof verification depends on the
648// HashFunction passed as parameter.
649func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool, error) {
650 siblings, err := UnpackSiblings(hashFunc, packedSiblings)
651 if err != nil {
652 return false, err
653 }
654
655 keyPath := make([]byte, hashFunc.Len())
656 copy(keyPath[:], k)
657
658 key, _, err := newLeafValue(hashFunc, k, v)
659 if err != nil {
660 return false, err
661 }
662
663 path := getPath(len(siblings), keyPath)
664 for i := len(siblings) - 1; i >= 0; i-- {
665 if path[i] {
666 key, _, err = newIntermediate(hashFunc, siblings[i], key)
667 if err != nil {
668 return false, err
669 }
670 } else {
671 key, _, err = newIntermediate(hashFunc, key, siblings[i])
672 if err != nil {
673 return false, err
674 }
675 }
676 }
677 if bytes.Equal(key[:], root) {
678 return true, nil
679 }
680 return false, nil
681}
682
683func (t *Tree) dbPut(k, v []byte) error {
684 if t.tx == nil {
685 return ErrDBNoTx
686 }
687 t.dbg.incDbPut()
688 return t.tx.Put(k, v)
689}
690
691func (t *Tree) dbGet(k []byte) ([]byte, error) {
692 // if key is empty, return empty as value
693 if bytes.Equal(k, t.emptyHash) {
694 return t.emptyHash, nil
695 }
696 t.dbg.incDbGet()
697
698 v, err := t.db.Get(k)
699 if err == nil {
700 return v, nil
701 }
702 if t.tx != nil {
703 return t.tx.Get(k)
704 }
705 return nil, db.ErrNotFound
706}
707
708// Warning: should be called with a Tree.tx created, and with a Tree.tx.Commit
709// after the setNLeafs call.
710func (t *Tree) incNLeafs(nLeafs int) error {
711 oldNLeafs, err := t.GetNLeafs()
712 if err != nil {
713 return err
714 }
715 newNLeafs := oldNLeafs + nLeafs
716 return t.setNLeafs(newNLeafs)
717}
718
719// Warning: should be called with a Tree.tx created, and with a Tree.tx.Commit
720// after the setNLeafs call.
721func (t *Tree) setNLeafs(nLeafs int) error {
722 b := make([]byte, 8)
723 binary.LittleEndian.PutUint64(b, uint64(nLeafs))
724 if err := t.dbPut(dbKeyNLeafs, b); err != nil {
725 return err
726 }
727 return nil
728}
729
730// GetNLeafs returns the number of Leafs of the Tree.
731func (t *Tree) GetNLeafs() (int, error) {
732 b, err := t.dbGet(dbKeyNLeafs)
733 if err != nil {
734 return 0, err
735 }
736 nLeafs := binary.LittleEndian.Uint64(b)
737 return int(nLeafs), nil
738}
739
740// Iterate iterates through the full Tree, executing the given function on each
741// node of the Tree.
742func (t *Tree) Iterate(rootKey []byte, f func([]byte, []byte)) error {
743 // allow to define which root to use
744 if rootKey == nil {
745 rootKey = t.Root()
746 }
747 return t.iter(rootKey, f)
748}
749
750// IterateWithStop does the same than Iterate, but with int for the current
751// level, and a boolean parameter used by the passed function, is to indicate to
752// stop iterating on the branch when the method returns 'true'.
753func (t *Tree) IterateWithStop(rootKey []byte, f func(int, []byte, []byte) bool) error {
754 // allow to define which root to use
755 if rootKey == nil {
756 rootKey = t.Root()
757 }
758 return t.iterWithStop(rootKey, 0, f)
759}
760
761func (t *Tree) iterWithStop(k []byte, currLevel int, f func(int, []byte, []byte) bool) error {
762 v, err := t.dbGet(k)
763 if err != nil {
764 return err
765 }
766 currLevel++
767
768 switch v[0] {
769 case PrefixValueEmpty:
770 f(currLevel, k, v)
771 case PrefixValueLeaf:
772 f(currLevel, k, v)
773 case PrefixValueIntermediate:
774 stop := f(currLevel, k, v)
775 if stop {
776 return nil
777 }
778 l, r := ReadIntermediateChilds(v)
779 if err = t.iterWithStop(l, currLevel, f); err != nil {
780 return err
781 }
782 if err = t.iterWithStop(r, currLevel, f); err != nil {
783 return err
784 }
785 default:
786 return ErrInvalidValuePrefix
787 }
788 return nil
789}
790
791func (t *Tree) iter(k []byte, f func([]byte, []byte)) error {
792 f2 := func(currLvl int, k, v []byte) bool {
793 f(k, v)
794 return false
795 }
796 return t.iterWithStop(k, 0, f2)
797}
798
799// Dump exports all the Tree leafs in a byte array of length:
800// [ N * (2+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
801// [ 1 byte | 1 byte | S bytes | len(v) bytes ]
802// [ len(k) | len(v) | key | value ]
803// Where S is the size of the output of the hash function used for the Tree.
804func (t *Tree) Dump(rootKey []byte) ([]byte, error) {
805 // allow to define which root to use
806 if rootKey == nil {
807 rootKey = t.Root()
808 }
809
810 // WARNING current encoding only supports key & values of 255 bytes each
811 // (due using only 1 byte for the length headers).
812 var b []byte
813 err := t.Iterate(rootKey, func(k, v []byte) {
814 if v[0] != PrefixValueLeaf {
815 return
816 }
817 leafK, leafV := ReadLeafValue(v)
818 kv := make([]byte, 2+len(leafK)+len(leafV))
819 kv[0] = byte(len(leafK))
820 kv[1] = byte(len(leafV))
821 copy(kv[2:2+len(leafK)], leafK)
822 copy(kv[2+len(leafK):], leafV)
823 b = append(b, kv...)
824 })
825 return b, err
826}
827
828// ImportDump imports the leafs (that have been exported with the ExportLeafs
829// method) in the Tree.
830func (t *Tree) ImportDump(b []byte) error {
831 r := bytes.NewReader(b)
832 var err error
833 var keys, values [][]byte
834 for {
835 l := make([]byte, 2)
836 _, err = io.ReadFull(r, l)
837 if err == io.EOF {
838 break
839 } else if err != nil {
840 return err
841 }
842 k := make([]byte, l[0])
843 _, err = io.ReadFull(r, k)
844 if err != nil {
845 return err
846 }
847 v := make([]byte, l[1])
848 _, err = io.ReadFull(r, v)
849 if err != nil {
850 return err
851 }
852 keys = append(keys, k)
853 values = append(values, v)
854 }
855 if _, err = t.AddBatch(keys, values); err != nil {
856 return err
857 }
858 return nil
859}
860
861// Graphviz iterates across the full tree to generate a string Graphviz
862// representation of the tree and writes it to w
863func (t *Tree) Graphviz(w io.Writer, rootKey []byte) error {
864 return t.GraphvizFirstNLevels(w, rootKey, t.maxLevels)
865}
866
867// GraphvizFirstNLevels iterates across the first NLevels of the tree to
868// generate a string Graphviz representation of the first NLevels of the tree
869// and writes it to w
870func (t *Tree) GraphvizFirstNLevels(w io.Writer, rootKey []byte, untilLvl int) error {
871 fmt.Fprintf(w, `digraph hierarchy {
872node [fontname=Monospace,fontsize=10,shape=box]
873`)
874 if rootKey == nil {
875 rootKey = t.Root()
876 }
877 nEmpties := 0
878 err := t.iterWithStop(rootKey, 0, func(currLvl int, k, v []byte) bool {
879 if currLvl == untilLvl {
880 return true // to stop the iter from going down
881 }
882 switch v[0] {
883 case PrefixValueEmpty:
884 case PrefixValueLeaf:
885 fmt.Fprintf(w, "\"%v\" [style=filled];\n", hex.EncodeToString(k[:nChars]))
886 // key & value from the leaf
887 kB, vB := ReadLeafValue(v)
888 fmt.Fprintf(w, "\"%v\" -> {\"k:%v\\nv:%v\"}\n",
889 hex.EncodeToString(k[:nChars]), hex.EncodeToString(kB[:nChars]),
890 hex.EncodeToString(vB[:nChars]))
891 fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
892 hex.EncodeToString(kB[:nChars]), hex.EncodeToString(vB[:nChars]))
893 case PrefixValueIntermediate:
894 l, r := ReadIntermediateChilds(v)
895 lStr := hex.EncodeToString(l[:nChars])
896 rStr := hex.EncodeToString(r[:nChars])
897 eStr := ""
898 if bytes.Equal(l, t.emptyHash) {
899 lStr = fmt.Sprintf("empty%v", nEmpties)
900 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
901 lStr)
902 nEmpties++
903 }
904 if bytes.Equal(r, t.emptyHash) {
905 rStr = fmt.Sprintf("empty%v", nEmpties)
906 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
907 rStr)
908 nEmpties++
909 }
910 fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", hex.EncodeToString(k[:nChars]),
911 lStr, rStr)
912 fmt.Fprint(w, eStr)
913 default:
914 }
915 return false
916 })
917 fmt.Fprintf(w, "}\n")
918 return err
919}
920
921// PrintGraphviz prints the output of Tree.Graphviz
922func (t *Tree) PrintGraphviz(rootKey []byte) error {
923 if rootKey == nil {
924 rootKey = t.Root()
925 }
926 return t.PrintGraphvizFirstNLevels(rootKey, t.maxLevels)
927}
928
929// PrintGraphvizFirstNLevels prints the output of Tree.GraphvizFirstNLevels
930func (t *Tree) PrintGraphvizFirstNLevels(rootKey []byte, untilLvl int) error {
931 if rootKey == nil {
932 rootKey = t.Root()
933 }
934 w := bytes.NewBufferString("")
935 fmt.Fprintf(w,
936 "--------\nGraphviz of the Tree with Root "+hex.EncodeToString(rootKey)+":\n")
937 err := t.GraphvizFirstNLevels(w, rootKey, untilLvl)
938 if err != nil {
939 fmt.Println(w)
940 return err
941 }
942 fmt.Fprintf(w,
943 "End of Graphviz of the Tree with Root "+hex.EncodeToString(rootKey)+"\n--------\n")
944
945 fmt.Println(w)
946 return nil
947}
948
949// Purge WIP: unimplemented TODO
950func (t *Tree) Purge(keys [][]byte) error {
951 return nil
952}
953
954// TODO circom proofs