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