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