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-1 {
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.GetNLeafs()
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 hashFunction: t.hashFunction,
864 dbg: t.dbg,
865 }, nil
866}
867
868// Iterate iterates through the full Tree, executing the given function on each
869// node of the Tree.
870func (t *Tree) Iterate(fromRoot []byte, f func([]byte, []byte)) error {
871 rTx := t.db.ReadTx()
872 defer rTx.Discard()
873
874 return t.IterateWithTx(rTx, fromRoot, f)
875}
876
877// IterateWithTx does the same than the Iterate method, but allowing to pass
878// the db.ReadTx that is used.
879func (t *Tree) IterateWithTx(rTx db.ReadTx, fromRoot []byte, f func([]byte, []byte)) error {
880 // allow to define which root to use
881 if fromRoot == nil {
882 var err error
883 fromRoot, err = t.RootWithTx(rTx)
884 if err != nil {
885 return err
886 }
887 }
888 return t.iter(rTx, fromRoot, f)
889}
890
891// IterateWithStop does the same than Iterate, but with int for the current
892// level, and a boolean parameter used by the passed function, is to indicate to
893// stop iterating on the branch when the method returns 'true'.
894func (t *Tree) IterateWithStop(fromRoot []byte, f func(int, []byte, []byte) bool) error {
895 rTx := t.db.ReadTx()
896 defer rTx.Discard()
897
898 // allow to define which root to use
899 if fromRoot == nil {
900 var err error
901 fromRoot, err = t.RootWithTx(rTx)
902 if err != nil {
903 return err
904 }
905 }
906 return t.iterWithStop(rTx, fromRoot, 0, f)
907}
908
909// IterateWithStopWithTx does the same than the IterateWithStop method, but
910// allowing to pass the db.ReadTx that is used.
911func (t *Tree) IterateWithStopWithTx(rTx db.ReadTx, fromRoot []byte,
912 f func(int, []byte, []byte) bool) error {
913 // allow to define which root to use
914 if fromRoot == nil {
915 var err error
916 fromRoot, err = t.RootWithTx(rTx)
917 if err != nil {
918 return err
919 }
920 }
921 return t.iterWithStop(rTx, fromRoot, 0, f)
922}
923
924func (t *Tree) iterWithStop(rTx db.ReadTx, k []byte, currLevel int,
925 f func(int, []byte, []byte) bool) error {
926 var v []byte
927 var err error
928 if bytes.Equal(k, t.emptyHash) {
929 v = t.emptyHash
930 } else {
931 v, err = rTx.Get(k)
932 if err != nil {
933 return err
934 }
935 }
936 currLevel++
937
938 switch v[0] {
939 case PrefixValueEmpty:
940 f(currLevel, k, v)
941 case PrefixValueLeaf:
942 f(currLevel, k, v)
943 case PrefixValueIntermediate:
944 stop := f(currLevel, k, v)
945 if stop {
946 return nil
947 }
948 l, r := ReadIntermediateChilds(v)
949 if err = t.iterWithStop(rTx, l, currLevel, f); err != nil {
950 return err
951 }
952 if err = t.iterWithStop(rTx, r, currLevel, f); err != nil {
953 return err
954 }
955 default:
956 return ErrInvalidValuePrefix
957 }
958 return nil
959}
960
961func (t *Tree) iter(rTx db.ReadTx, k []byte, f func([]byte, []byte)) error {
962 f2 := func(currLvl int, k, v []byte) bool {
963 f(k, v)
964 return false
965 }
966 return t.iterWithStop(rTx, k, 0, f2)
967}
968
969// Dump exports all the Tree leafs in a byte array of length:
970// [ N * (2+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
971// [ 1 byte | 1 byte | S bytes | len(v) bytes ]
972// [ len(k) | len(v) | key | value ]
973// Where S is the size of the output of the hash function used for the Tree.
974func (t *Tree) Dump(fromRoot []byte) ([]byte, error) {
975 // allow to define which root to use
976 if fromRoot == nil {
977 var err error
978 fromRoot, err = t.Root()
979 if err != nil {
980 return nil, err
981 }
982 }
983
984 // WARNING current encoding only supports key & values of 255 bytes each
985 // (due using only 1 byte for the length headers).
986 var b []byte
987 err := t.Iterate(fromRoot, func(k, v []byte) {
988 if v[0] != PrefixValueLeaf {
989 return
990 }
991 leafK, leafV := ReadLeafValue(v)
992 kv := make([]byte, 2+len(leafK)+len(leafV))
993 kv[0] = byte(len(leafK))
994 kv[1] = byte(len(leafV))
995 copy(kv[2:2+len(leafK)], leafK)
996 copy(kv[2+len(leafK):], leafV)
997 b = append(b, kv...)
998 })
999 return b, err
1000}
1001
1002// ImportDump imports the leafs (that have been exported with the Dump method)
1003// in the Tree.
1004func (t *Tree) ImportDump(b []byte) error {
1005 if !t.editable() {
1006 return ErrSnapshotNotEditable
1007 }
1008 root, err := t.Root()
1009 if err != nil {
1010 return err
1011 }
1012 if !bytes.Equal(root, t.emptyHash) {
1013 return ErrTreeNotEmpty
1014 }
1015
1016 r := bytes.NewReader(b)
1017 var keys, values [][]byte
1018 for {
1019 l := make([]byte, 2)
1020 _, err = io.ReadFull(r, l)
1021 if err == io.EOF {
1022 break
1023 } else if err != nil {
1024 return err
1025 }
1026 k := make([]byte, l[0])
1027 _, err = io.ReadFull(r, k)
1028 if err != nil {
1029 return err
1030 }
1031 v := make([]byte, l[1])
1032 _, err = io.ReadFull(r, v)
1033 if err != nil {
1034 return err
1035 }
1036 keys = append(keys, k)
1037 values = append(values, v)
1038 }
1039 if _, err = t.AddBatch(keys, values); err != nil {
1040 return err
1041 }
1042 return nil
1043}
1044
1045// Graphviz iterates across the full tree to generate a string Graphviz
1046// representation of the tree and writes it to w
1047func (t *Tree) Graphviz(w io.Writer, fromRoot []byte) error {
1048 return t.GraphvizFirstNLevels(w, fromRoot, t.maxLevels)
1049}
1050
1051// GraphvizFirstNLevels iterates across the first NLevels of the tree to
1052// generate a string Graphviz representation of the first NLevels of the tree
1053// and writes it to w
1054func (t *Tree) GraphvizFirstNLevels(w io.Writer, fromRoot []byte, untilLvl int) error {
1055 fmt.Fprintf(w, `digraph hierarchy {
1056node [fontname=Monospace,fontsize=10,shape=box]
1057`)
1058
1059 rTx := t.db.ReadTx()
1060 defer rTx.Discard()
1061
1062 if fromRoot == nil {
1063 var err error
1064 fromRoot, err = t.RootWithTx(rTx)
1065 if err != nil {
1066 return err
1067 }
1068 }
1069
1070 nEmpties := 0
1071 err := t.iterWithStop(rTx, fromRoot, 0, func(currLvl int, k, v []byte) bool {
1072 if currLvl == untilLvl {
1073 return true // to stop the iter from going down
1074 }
1075 switch v[0] {
1076 case PrefixValueEmpty:
1077 case PrefixValueLeaf:
1078 fmt.Fprintf(w, "\"%v\" [style=filled];\n", hex.EncodeToString(k[:nChars]))
1079 // key & value from the leaf
1080 kB, vB := ReadLeafValue(v)
1081 fmt.Fprintf(w, "\"%v\" -> {\"k:%v\\nv:%v\"}\n",
1082 hex.EncodeToString(k[:nChars]), hex.EncodeToString(kB[:nChars]),
1083 hex.EncodeToString(vB[:nChars]))
1084 fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
1085 hex.EncodeToString(kB[:nChars]), hex.EncodeToString(vB[:nChars]))
1086 case PrefixValueIntermediate:
1087 l, r := ReadIntermediateChilds(v)
1088 lStr := hex.EncodeToString(l[:nChars])
1089 rStr := hex.EncodeToString(r[:nChars])
1090 eStr := ""
1091 if bytes.Equal(l, t.emptyHash) {
1092 lStr = fmt.Sprintf("empty%v", nEmpties)
1093 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
1094 lStr)
1095 nEmpties++
1096 }
1097 if bytes.Equal(r, t.emptyHash) {
1098 rStr = fmt.Sprintf("empty%v", nEmpties)
1099 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
1100 rStr)
1101 nEmpties++
1102 }
1103 fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", hex.EncodeToString(k[:nChars]),
1104 lStr, rStr)
1105 fmt.Fprint(w, eStr)
1106 default:
1107 }
1108 return false
1109 })
1110 fmt.Fprintf(w, "}\n")
1111 return err
1112}
1113
1114// PrintGraphviz prints the output of Tree.Graphviz
1115func (t *Tree) PrintGraphviz(fromRoot []byte) error {
1116 if fromRoot == nil {
1117 var err error
1118 fromRoot, err = t.Root()
1119 if err != nil {
1120 return err
1121 }
1122 }
1123 return t.PrintGraphvizFirstNLevels(fromRoot, t.maxLevels)
1124}
1125
1126// PrintGraphvizFirstNLevels prints the output of Tree.GraphvizFirstNLevels
1127func (t *Tree) PrintGraphvizFirstNLevels(fromRoot []byte, untilLvl int) error {
1128 if fromRoot == nil {
1129 var err error
1130 fromRoot, err = t.Root()
1131 if err != nil {
1132 return err
1133 }
1134 }
1135 w := bytes.NewBufferString("")
1136 fmt.Fprintf(w,
1137 "--------\nGraphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+":\n")
1138 err := t.GraphvizFirstNLevels(w, fromRoot, untilLvl)
1139 if err != nil {
1140 fmt.Println(w)
1141 return err
1142 }
1143 fmt.Fprintf(w,
1144 "End of Graphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+"\n--------\n")
1145
1146 fmt.Println(w)
1147 return nil
1148}
1149
1150// TODO circom proofs
1151// TODO data structure for proofs (including root, key, value, siblings,
1152// hashFunction) + method to verify that data structure