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 leafV, 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 var err error
715 var keys, values [][]byte
716 for {
717 l := make([]byte, 2)
718 _, err = io.ReadFull(r, l)
719 if err == io.EOF {
720 break
721 } else if err != nil {
722 return err
723 }
724 k := make([]byte, l[0])
725 _, err = io.ReadFull(r, k)
726 if err != nil {
727 return err
728 }
729 v := make([]byte, l[1])
730 _, err = io.ReadFull(r, v)
731 if err != nil {
732 return err
733 }
734 keys = append(keys, k)
735 values = append(values, v)
736 }
737 if _, err = t.AddBatch(keys, values); err != nil {
738 return err
739 }
740 return nil
741}
742
743// Graphviz iterates across the full tree to generate a string Graphviz
744// representation of the tree and writes it to w
745func (t *Tree) Graphviz(w io.Writer, rootKey []byte) error {
746 return t.GraphvizFirstNLevels(w, rootKey, t.maxLevels)
747}
748
749// GraphvizFirstNLevels iterates across the first NLevels of the tree to
750// generate a string Graphviz representation of the first NLevels of the tree
751// and writes it to w
752func (t *Tree) GraphvizFirstNLevels(w io.Writer, rootKey []byte, untilLvl int) error {
753 fmt.Fprintf(w, `digraph hierarchy {
754node [fontname=Monospace,fontsize=10,shape=box]
755`)
756 nChars := 4 // TODO move to global constant
757 nEmpties := 0
758 err := t.iterWithStop(t.root, 0, func(currLvl int, k, v []byte) bool {
759 if currLvl == untilLvl {
760 return true // to stop the iter from going down
761 }
762 switch v[0] {
763 case PrefixValueEmpty:
764 case PrefixValueLeaf:
765 fmt.Fprintf(w, "\"%v\" [style=filled];\n", hex.EncodeToString(k[:nChars]))
766 // key & value from the leaf
767 kB, vB := ReadLeafValue(v)
768 fmt.Fprintf(w, "\"%v\" -> {\"k:%v\\nv:%v\"}\n",
769 hex.EncodeToString(k[:nChars]), hex.EncodeToString(kB[:nChars]),
770 hex.EncodeToString(vB[:nChars]))
771 fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
772 hex.EncodeToString(kB[:nChars]), hex.EncodeToString(vB[:nChars]))
773 case PrefixValueIntermediate:
774 l, r := ReadIntermediateChilds(v)
775 lStr := hex.EncodeToString(l[:nChars])
776 rStr := hex.EncodeToString(r[:nChars])
777 eStr := ""
778 if bytes.Equal(l, t.emptyHash) {
779 lStr = fmt.Sprintf("empty%v", nEmpties)
780 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
781 lStr)
782 nEmpties++
783 }
784 if bytes.Equal(r, t.emptyHash) {
785 rStr = fmt.Sprintf("empty%v", nEmpties)
786 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
787 rStr)
788 nEmpties++
789 }
790 fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", hex.EncodeToString(k[:nChars]),
791 lStr, rStr)
792 fmt.Fprint(w, eStr)
793 default:
794 }
795 return false
796 })
797 fmt.Fprintf(w, "}\n")
798 return err
799}
800
801// PrintGraphviz prints the output of Tree.Graphviz
802func (t *Tree) PrintGraphviz(rootKey []byte) error {
803 return t.PrintGraphvizFirstNLevels(rootKey, t.maxLevels)
804}
805
806// PrintGraphvizFirstNLevels prints the output of Tree.GraphvizFirstNLevels
807func (t *Tree) PrintGraphvizFirstNLevels(rootKey []byte, untilLvl int) error {
808 if rootKey == nil {
809 rootKey = t.Root()
810 }
811 w := bytes.NewBufferString("")
812 fmt.Fprintf(w,
813 "--------\nGraphviz of the Tree with Root "+hex.EncodeToString(rootKey)+":\n")
814 err := t.GraphvizFirstNLevels(w, nil, untilLvl)
815 if err != nil {
816 fmt.Println(w)
817 return err
818 }
819 fmt.Fprintf(w,
820 "End of Graphviz of the Tree with Root "+hex.EncodeToString(rootKey)+"\n--------\n")
821
822 fmt.Println(w)
823 return nil
824}
825
826// Purge WIP: unimplemented
827func (t *Tree) Purge(keys [][]byte) error {
828 return nil
829}