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