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// Iterate iterates through the full Tree, executing the given function on each
760// node of the Tree.
761func (t *Tree) Iterate(rootKey []byte, f func([]byte, []byte)) error {
762 // allow to define which root to use
763 if rootKey == nil {
764 rootKey = t.Root()
765 }
766 return t.iter(rootKey, f)
767}
768
769// IterateWithStop does the same than Iterate, but with int for the current
770// level, and a boolean parameter used by the passed function, is to indicate to
771// stop iterating on the branch when the method returns 'true'.
772func (t *Tree) IterateWithStop(rootKey []byte, f func(int, []byte, []byte) bool) error {
773 // allow to define which root to use
774 if rootKey == nil {
775 rootKey = t.Root()
776 }
777 return t.iterWithStop(rootKey, 0, f)
778}
779
780func (t *Tree) iterWithStop(k []byte, currLevel int, f func(int, []byte, []byte) bool) error {
781 v, err := t.dbGet(k)
782 if err != nil {
783 return err
784 }
785 currLevel++
786
787 switch v[0] {
788 case PrefixValueEmpty:
789 f(currLevel, k, v)
790 case PrefixValueLeaf:
791 f(currLevel, k, v)
792 case PrefixValueIntermediate:
793 stop := f(currLevel, k, v)
794 if stop {
795 return nil
796 }
797 l, r := ReadIntermediateChilds(v)
798 if err = t.iterWithStop(l, currLevel, f); err != nil {
799 return err
800 }
801 if err = t.iterWithStop(r, currLevel, f); err != nil {
802 return err
803 }
804 default:
805 return ErrInvalidValuePrefix
806 }
807 return nil
808}
809
810func (t *Tree) iter(k []byte, f func([]byte, []byte)) error {
811 f2 := func(currLvl int, k, v []byte) bool {
812 f(k, v)
813 return false
814 }
815 return t.iterWithStop(k, 0, f2)
816}
817
818// Dump exports all the Tree leafs in a byte array of length:
819// [ N * (2+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
820// [ 1 byte | 1 byte | S bytes | len(v) bytes ]
821// [ len(k) | len(v) | key | value ]
822// Where S is the size of the output of the hash function used for the Tree.
823func (t *Tree) Dump(rootKey []byte) ([]byte, error) {
824 // allow to define which root to use
825 if rootKey == nil {
826 rootKey = t.Root()
827 }
828
829 // WARNING current encoding only supports key & values of 255 bytes each
830 // (due using only 1 byte for the length headers).
831 var b []byte
832 err := t.Iterate(rootKey, func(k, v []byte) {
833 if v[0] != PrefixValueLeaf {
834 return
835 }
836 leafK, leafV := ReadLeafValue(v)
837 kv := make([]byte, 2+len(leafK)+len(leafV))
838 kv[0] = byte(len(leafK))
839 kv[1] = byte(len(leafV))
840 copy(kv[2:2+len(leafK)], leafK)
841 copy(kv[2+len(leafK):], leafV)
842 b = append(b, kv...)
843 })
844 return b, err
845}
846
847// ImportDump imports the leafs (that have been exported with the ExportLeafs
848// method) in the Tree.
849func (t *Tree) ImportDump(b []byte) error {
850 r := bytes.NewReader(b)
851 var err error
852 var keys, values [][]byte
853 for {
854 l := make([]byte, 2)
855 _, err = io.ReadFull(r, l)
856 if err == io.EOF {
857 break
858 } else if err != nil {
859 return err
860 }
861 k := make([]byte, l[0])
862 _, err = io.ReadFull(r, k)
863 if err != nil {
864 return err
865 }
866 v := make([]byte, l[1])
867 _, err = io.ReadFull(r, v)
868 if err != nil {
869 return err
870 }
871 keys = append(keys, k)
872 values = append(values, v)
873 }
874 if _, err = t.AddBatch(keys, values); err != nil {
875 return err
876 }
877 return nil
878}
879
880// Graphviz iterates across the full tree to generate a string Graphviz
881// representation of the tree and writes it to w
882func (t *Tree) Graphviz(w io.Writer, rootKey []byte) error {
883 return t.GraphvizFirstNLevels(w, rootKey, t.maxLevels)
884}
885
886// GraphvizFirstNLevels iterates across the first NLevels of the tree to
887// generate a string Graphviz representation of the first NLevels of the tree
888// and writes it to w
889func (t *Tree) GraphvizFirstNLevels(w io.Writer, rootKey []byte, untilLvl int) error {
890 fmt.Fprintf(w, `digraph hierarchy {
891node [fontname=Monospace,fontsize=10,shape=box]
892`)
893 if rootKey == nil {
894 rootKey = t.Root()
895 }
896 nEmpties := 0
897 err := t.iterWithStop(rootKey, 0, func(currLvl int, k, v []byte) bool {
898 if currLvl == untilLvl {
899 return true // to stop the iter from going down
900 }
901 switch v[0] {
902 case PrefixValueEmpty:
903 case PrefixValueLeaf:
904 fmt.Fprintf(w, "\"%v\" [style=filled];\n", hex.EncodeToString(k[:nChars]))
905 // key & value from the leaf
906 kB, vB := ReadLeafValue(v)
907 fmt.Fprintf(w, "\"%v\" -> {\"k:%v\\nv:%v\"}\n",
908 hex.EncodeToString(k[:nChars]), hex.EncodeToString(kB[:nChars]),
909 hex.EncodeToString(vB[:nChars]))
910 fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
911 hex.EncodeToString(kB[:nChars]), hex.EncodeToString(vB[:nChars]))
912 case PrefixValueIntermediate:
913 l, r := ReadIntermediateChilds(v)
914 lStr := hex.EncodeToString(l[:nChars])
915 rStr := hex.EncodeToString(r[:nChars])
916 eStr := ""
917 if bytes.Equal(l, t.emptyHash) {
918 lStr = fmt.Sprintf("empty%v", nEmpties)
919 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
920 lStr)
921 nEmpties++
922 }
923 if bytes.Equal(r, t.emptyHash) {
924 rStr = fmt.Sprintf("empty%v", nEmpties)
925 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
926 rStr)
927 nEmpties++
928 }
929 fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", hex.EncodeToString(k[:nChars]),
930 lStr, rStr)
931 fmt.Fprint(w, eStr)
932 default:
933 }
934 return false
935 })
936 fmt.Fprintf(w, "}\n")
937 return err
938}
939
940// PrintGraphviz prints the output of Tree.Graphviz
941func (t *Tree) PrintGraphviz(rootKey []byte) error {
942 if rootKey == nil {
943 rootKey = t.Root()
944 }
945 return t.PrintGraphvizFirstNLevels(rootKey, t.maxLevels)
946}
947
948// PrintGraphvizFirstNLevels prints the output of Tree.GraphvizFirstNLevels
949func (t *Tree) PrintGraphvizFirstNLevels(rootKey []byte, untilLvl int) error {
950 if rootKey == nil {
951 rootKey = t.Root()
952 }
953 w := bytes.NewBufferString("")
954 fmt.Fprintf(w,
955 "--------\nGraphviz of the Tree with Root "+hex.EncodeToString(rootKey)+":\n")
956 err := t.GraphvizFirstNLevels(w, rootKey, untilLvl)
957 if err != nil {
958 fmt.Println(w)
959 return err
960 }
961 fmt.Fprintf(w,
962 "End of Graphviz of the Tree with Root "+hex.EncodeToString(rootKey)+"\n--------\n")
963
964 fmt.Println(w)
965 return nil
966}
967
968// TODO circom proofs
969// TODO data structure for proofs (including root, key, value, siblings,
970// hashFunction) + method to verify that data structure