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