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