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