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