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