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. If the key exists in
538// the Tree, the proof will be of existence, if the key does not exist in the
539// tree, the proof will be of non-existence.
540func (t *Tree) GenProof(k []byte) ([]byte, []byte, 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, err
550 }
551
552 leafK, leafV := ReadLeafValue(value)
553 if !bytes.Equal(k, leafK) {
554 fmt.Println("key not in Tree")
555 fmt.Println(leafK)
556 fmt.Println(leafV)
557 // TODO proof of non-existence
558 panic("unimplemented")
559 }
560
561 s := PackSiblings(t.hashFunction, siblings)
562 return leafV, s, nil
563}
564
565// PackSiblings packs the siblings into a byte array.
566// [ 1 byte | L bytes | S * N bytes ]
567// [ bitmap length (L) | bitmap | N non-zero siblings ]
568// Where the bitmap indicates if the sibling is 0 or a value from the siblings
569// array. And S is the size of the output of the hash function used for the
570// Tree.
571func PackSiblings(hashFunc HashFunction, siblings [][]byte) []byte {
572 var b []byte
573 var bitmap []bool
574 emptySibling := make([]byte, hashFunc.Len())
575 for i := 0; i < len(siblings); i++ {
576 if bytes.Equal(siblings[i], emptySibling) {
577 bitmap = append(bitmap, false)
578 } else {
579 bitmap = append(bitmap, true)
580 b = append(b, siblings[i]...)
581 }
582 }
583
584 bitmapBytes := bitmapToBytes(bitmap)
585 l := len(bitmapBytes)
586
587 res := make([]byte, l+1+len(b))
588 res[0] = byte(l) // set the bitmapBytes length
589 copy(res[1:1+l], bitmapBytes)
590 copy(res[1+l:], b)
591 return res
592}
593
594// UnpackSiblings unpacks the siblings from a byte array.
595func UnpackSiblings(hashFunc HashFunction, b []byte) ([][]byte, error) {
596 l := b[0]
597 bitmapBytes := b[1 : 1+l]
598 bitmap := bytesToBitmap(bitmapBytes)
599 siblingsBytes := b[1+l:]
600 iSibl := 0
601 emptySibl := make([]byte, hashFunc.Len())
602 var siblings [][]byte
603 for i := 0; i < len(bitmap); i++ {
604 if iSibl >= len(siblingsBytes) {
605 break
606 }
607 if bitmap[i] {
608 siblings = append(siblings, siblingsBytes[iSibl:iSibl+hashFunc.Len()])
609 iSibl += hashFunc.Len()
610 } else {
611 siblings = append(siblings, emptySibl)
612 }
613 }
614 return siblings, nil
615}
616
617func bitmapToBytes(bitmap []bool) []byte {
618 bitmapBytesLen := int(math.Ceil(float64(len(bitmap)) / 8)) //nolint:gomnd
619 b := make([]byte, bitmapBytesLen)
620 for i := 0; i < len(bitmap); i++ {
621 if bitmap[i] {
622 b[i/8] |= 1 << (i % 8)
623 }
624 }
625 return b
626}
627
628func bytesToBitmap(b []byte) []bool {
629 var bitmap []bool
630 for i := 0; i < len(b); i++ {
631 for j := 0; j < 8; j++ {
632 bitmap = append(bitmap, b[i]&(1<<j) > 0)
633 }
634 }
635 return bitmap
636}
637
638// Get returns the value for a given key
639func (t *Tree) Get(k []byte) ([]byte, []byte, error) {
640 keyPath := make([]byte, t.hashFunction.Len())
641 copy(keyPath[:], k)
642
643 path := getPath(t.maxLevels, keyPath)
644 // go down to the leaf
645 var siblings [][]byte
646 _, value, _, err := t.down(k, t.root, siblings, path, 0, true)
647 if err != nil {
648 return nil, nil, err
649 }
650 leafK, leafV := ReadLeafValue(value)
651 if !bytes.Equal(k, leafK) {
652 return leafK, leafV, fmt.Errorf("Tree.Get error: keys doesn't match, %s != %s",
653 BytesToBigInt(k), BytesToBigInt(leafK))
654 }
655
656 return leafK, leafV, nil
657}
658
659// CheckProof verifies the given proof. The proof verification depends on the
660// HashFunction passed as parameter.
661func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool, error) {
662 siblings, err := UnpackSiblings(hashFunc, packedSiblings)
663 if err != nil {
664 return false, err
665 }
666
667 keyPath := make([]byte, hashFunc.Len())
668 copy(keyPath[:], k)
669
670 key, _, err := newLeafValue(hashFunc, k, v)
671 if err != nil {
672 return false, err
673 }
674
675 path := getPath(len(siblings), keyPath)
676 for i := len(siblings) - 1; i >= 0; i-- {
677 if path[i] {
678 key, _, err = newIntermediate(hashFunc, siblings[i], key)
679 if err != nil {
680 return false, err
681 }
682 } else {
683 key, _, err = newIntermediate(hashFunc, key, siblings[i])
684 if err != nil {
685 return false, err
686 }
687 }
688 }
689 if bytes.Equal(key[:], root) {
690 return true, nil
691 }
692 return false, nil
693}
694
695func (t *Tree) dbPut(k, v []byte) error {
696 if t.dbBatch == nil {
697 return ErrDBNoTx
698 }
699 t.dbg.incDbPut()
700 t.batchMemory.Put(k, v) // TODO TMP
701 return t.dbBatch.Put(k, v)
702}
703
704func (t *Tree) dbGet(k []byte) ([]byte, error) {
705 // if key is empty, return empty as value
706 if bytes.Equal(k, t.emptyHash) {
707 return t.emptyHash, nil
708 }
709 t.dbg.incDbGet()
710
711 v, err := t.db.Get(k)
712 if err == nil {
713 return v, nil
714 }
715 if t.dbBatch != nil {
716 // TODO TMP
717 v, ok := t.batchMemory.Get(k)
718 if !ok {
719 return nil, ErrKeyNotFound
720 }
721 // /TMP
722 return v, nil
723 }
724 return nil, ErrKeyNotFound
725}
726
727// Warning: should be called with a Tree.dbBatch created, and with a
728// Tree.dbBatch.Write after the incNLeafs call.
729func (t *Tree) incNLeafs(nLeafs int) error {
730 oldNLeafs, err := t.GetNLeafs()
731 if err != nil {
732 return err
733 }
734 newNLeafs := oldNLeafs + nLeafs
735 return t.setNLeafs(newNLeafs)
736}
737
738// Warning: should be called with a Tree.dbBatch created, and with a
739// Tree.dbBatch.Write after the setNLeafs call.
740func (t *Tree) setNLeafs(nLeafs int) error {
741 b := make([]byte, 8)
742 binary.LittleEndian.PutUint64(b, uint64(nLeafs))
743 if err := t.dbPut(dbKeyNLeafs, b); err != nil {
744 return err
745 }
746 return nil
747}
748
749// GetNLeafs returns the number of Leafs of the Tree.
750func (t *Tree) GetNLeafs() (int, error) {
751 b, err := t.dbGet(dbKeyNLeafs)
752 if err != nil {
753 return 0, err
754 }
755 nLeafs := binary.LittleEndian.Uint64(b)
756 return int(nLeafs), nil
757}
758
759// Snapshot returns a copy of the Tree from the given root
760func (t *Tree) Snapshot(rootKey []byte) (*Tree, error) {
761 // TODO currently this method only changes the 'root pointer', but the
762 // db continues being the same. In a future iteration, once the
763 // db.Database interface allows to do database checkpoints, this method
764 // could be updated to do a full checkpoint of the database for the
765 // snapshot, to return a completly new independent tree containing the
766 // snapshot.
767 t.RLock()
768 defer t.RUnlock()
769
770 // allow to define which root to use
771 if rootKey == nil {
772 rootKey = t.Root()
773 }
774 return &Tree{
775 db: t.db,
776 maxLevels: t.maxLevels,
777 root: rootKey,
778 hashFunction: t.hashFunction,
779 dbg: t.dbg,
780 }, nil
781}
782
783// Iterate iterates through the full Tree, executing the given function on each
784// node of the Tree.
785func (t *Tree) Iterate(rootKey []byte, f func([]byte, []byte)) error {
786 // allow to define which root to use
787 if rootKey == nil {
788 rootKey = t.Root()
789 }
790 return t.iter(rootKey, f)
791}
792
793// IterateWithStop does the same than Iterate, but with int for the current
794// level, and a boolean parameter used by the passed function, is to indicate to
795// stop iterating on the branch when the method returns 'true'.
796func (t *Tree) IterateWithStop(rootKey []byte, f func(int, []byte, []byte) bool) error {
797 // allow to define which root to use
798 if rootKey == nil {
799 rootKey = t.Root()
800 }
801 return t.iterWithStop(rootKey, 0, f)
802}
803
804func (t *Tree) iterWithStop(k []byte, currLevel int, f func(int, []byte, []byte) bool) error {
805 v, err := t.dbGet(k)
806 if err != nil {
807 return err
808 }
809 currLevel++
810
811 switch v[0] {
812 case PrefixValueEmpty:
813 f(currLevel, k, v)
814 case PrefixValueLeaf:
815 f(currLevel, k, v)
816 case PrefixValueIntermediate:
817 stop := f(currLevel, k, v)
818 if stop {
819 return nil
820 }
821 l, r := ReadIntermediateChilds(v)
822 if err = t.iterWithStop(l, currLevel, f); err != nil {
823 return err
824 }
825 if err = t.iterWithStop(r, currLevel, f); err != nil {
826 return err
827 }
828 default:
829 return ErrInvalidValuePrefix
830 }
831 return nil
832}
833
834func (t *Tree) iter(k []byte, f func([]byte, []byte)) error {
835 f2 := func(currLvl int, k, v []byte) bool {
836 f(k, v)
837 return false
838 }
839 return t.iterWithStop(k, 0, f2)
840}
841
842// Dump exports all the Tree leafs in a byte array of length:
843// [ N * (2+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
844// [ 1 byte | 1 byte | S bytes | len(v) bytes ]
845// [ len(k) | len(v) | key | value ]
846// Where S is the size of the output of the hash function used for the Tree.
847func (t *Tree) Dump(rootKey []byte) ([]byte, error) {
848 // allow to define which root to use
849 if rootKey == nil {
850 rootKey = t.Root()
851 }
852
853 // WARNING current encoding only supports key & values of 255 bytes each
854 // (due using only 1 byte for the length headers).
855 var b []byte
856 err := t.Iterate(rootKey, func(k, v []byte) {
857 if v[0] != PrefixValueLeaf {
858 return
859 }
860 leafK, leafV := ReadLeafValue(v)
861 kv := make([]byte, 2+len(leafK)+len(leafV))
862 kv[0] = byte(len(leafK))
863 kv[1] = byte(len(leafV))
864 copy(kv[2:2+len(leafK)], leafK)
865 copy(kv[2+len(leafK):], leafV)
866 b = append(b, kv...)
867 })
868 return b, err
869}
870
871// ImportDump imports the leafs (that have been exported with the Dump method)
872// in the Tree.
873func (t *Tree) ImportDump(b []byte) error {
874 r := bytes.NewReader(b)
875 var err error
876 var keys, values [][]byte
877 for {
878 l := make([]byte, 2)
879 _, err = io.ReadFull(r, l)
880 if err == io.EOF {
881 break
882 } else if err != nil {
883 return err
884 }
885 k := make([]byte, l[0])
886 _, err = io.ReadFull(r, k)
887 if err != nil {
888 return err
889 }
890 v := make([]byte, l[1])
891 _, err = io.ReadFull(r, v)
892 if err != nil {
893 return err
894 }
895 keys = append(keys, k)
896 values = append(values, v)
897 }
898 if _, err = t.AddBatch(keys, values); err != nil {
899 return err
900 }
901 return nil
902}
903
904// Graphviz iterates across the full tree to generate a string Graphviz
905// representation of the tree and writes it to w
906func (t *Tree) Graphviz(w io.Writer, rootKey []byte) error {
907 return t.GraphvizFirstNLevels(w, rootKey, t.maxLevels)
908}
909
910// GraphvizFirstNLevels iterates across the first NLevels of the tree to
911// generate a string Graphviz representation of the first NLevels of the tree
912// and writes it to w
913func (t *Tree) GraphvizFirstNLevels(w io.Writer, rootKey []byte, untilLvl int) error {
914 fmt.Fprintf(w, `digraph hierarchy {
915node [fontname=Monospace,fontsize=10,shape=box]
916`)
917 if rootKey == nil {
918 rootKey = t.Root()
919 }
920 nEmpties := 0
921 err := t.iterWithStop(rootKey, 0, func(currLvl int, k, v []byte) bool {
922 if currLvl == untilLvl {
923 return true // to stop the iter from going down
924 }
925 switch v[0] {
926 case PrefixValueEmpty:
927 case PrefixValueLeaf:
928 fmt.Fprintf(w, "\"%v\" [style=filled];\n", hex.EncodeToString(k[:nChars]))
929 // key & value from the leaf
930 kB, vB := ReadLeafValue(v)
931 fmt.Fprintf(w, "\"%v\" -> {\"k:%v\\nv:%v\"}\n",
932 hex.EncodeToString(k[:nChars]), hex.EncodeToString(kB[:nChars]),
933 hex.EncodeToString(vB[:nChars]))
934 fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
935 hex.EncodeToString(kB[:nChars]), hex.EncodeToString(vB[:nChars]))
936 case PrefixValueIntermediate:
937 l, r := ReadIntermediateChilds(v)
938 lStr := hex.EncodeToString(l[:nChars])
939 rStr := hex.EncodeToString(r[:nChars])
940 eStr := ""
941 if bytes.Equal(l, t.emptyHash) {
942 lStr = fmt.Sprintf("empty%v", nEmpties)
943 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
944 lStr)
945 nEmpties++
946 }
947 if bytes.Equal(r, t.emptyHash) {
948 rStr = fmt.Sprintf("empty%v", nEmpties)
949 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
950 rStr)
951 nEmpties++
952 }
953 fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", hex.EncodeToString(k[:nChars]),
954 lStr, rStr)
955 fmt.Fprint(w, eStr)
956 default:
957 }
958 return false
959 })
960 fmt.Fprintf(w, "}\n")
961 return err
962}
963
964// PrintGraphviz prints the output of Tree.Graphviz
965func (t *Tree) PrintGraphviz(rootKey []byte) error {
966 if rootKey == nil {
967 rootKey = t.Root()
968 }
969 return t.PrintGraphvizFirstNLevels(rootKey, t.maxLevels)
970}
971
972// PrintGraphvizFirstNLevels prints the output of Tree.GraphvizFirstNLevels
973func (t *Tree) PrintGraphvizFirstNLevels(rootKey []byte, untilLvl int) error {
974 if rootKey == nil {
975 rootKey = t.Root()
976 }
977 w := bytes.NewBufferString("")
978 fmt.Fprintf(w,
979 "--------\nGraphviz of the Tree with Root "+hex.EncodeToString(rootKey)+":\n")
980 err := t.GraphvizFirstNLevels(w, rootKey, untilLvl)
981 if err != nil {
982 fmt.Println(w)
983 return err
984 }
985 fmt.Fprintf(w,
986 "End of Graphviz of the Tree with Root "+hex.EncodeToString(rootKey)+"\n--------\n")
987
988 fmt.Println(w)
989 return nil
990}
991
992// TODO circom proofs
993// TODO data structure for proofs (including root, key, value, siblings,
994// hashFunction) + method to verify that data structure