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