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.Mutex
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// [ 2 byte | 2 byte | L bytes | S * N bytes ]
663// [ full length | 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. The 2 2-byte that define the full length and bitmap length, are
667// encoded in little-endian.
668func PackSiblings(hashFunc HashFunction, siblings [][]byte) []byte {
669 var b []byte
670 var bitmap []bool
671 emptySibling := make([]byte, hashFunc.Len())
672 for i := 0; i < len(siblings); i++ {
673 if bytes.Equal(siblings[i], emptySibling) {
674 bitmap = append(bitmap, false)
675 } else {
676 bitmap = append(bitmap, true)
677 b = append(b, siblings[i]...)
678 }
679 }
680
681 bitmapBytes := bitmapToBytes(bitmap)
682 l := len(bitmapBytes)
683
684 fullLen := 4 + l + len(b) //nolint:gomnd
685 res := make([]byte, fullLen)
686 binary.LittleEndian.PutUint16(res[0:2], uint16(fullLen)) // set full length
687 binary.LittleEndian.PutUint16(res[2:4], uint16(l)) // set the bitmapBytes length
688 copy(res[4:4+l], bitmapBytes)
689 copy(res[4+l:], b)
690 return res
691}
692
693// UnpackSiblings unpacks the siblings from a byte array.
694func UnpackSiblings(hashFunc HashFunction, b []byte) ([][]byte, error) {
695 fullLen := binary.LittleEndian.Uint16(b[0:2])
696 l := binary.LittleEndian.Uint16(b[2:4]) // bitmap bytes length
697 if len(b) != int(fullLen) {
698 return nil,
699 fmt.Errorf("error unpacking siblings. Expected len: %d, current len: %d",
700 fullLen, len(b))
701 }
702
703 bitmapBytes := b[4 : 4+l]
704 bitmap := bytesToBitmap(bitmapBytes)
705 siblingsBytes := b[4+l:]
706 iSibl := 0
707 emptySibl := make([]byte, hashFunc.Len())
708 var siblings [][]byte
709 for i := 0; i < len(bitmap); i++ {
710 if iSibl >= len(siblingsBytes) {
711 break
712 }
713 if bitmap[i] {
714 siblings = append(siblings, siblingsBytes[iSibl:iSibl+hashFunc.Len()])
715 iSibl += hashFunc.Len()
716 } else {
717 siblings = append(siblings, emptySibl)
718 }
719 }
720 return siblings, nil
721}
722
723func bitmapToBytes(bitmap []bool) []byte {
724 bitmapBytesLen := int(math.Ceil(float64(len(bitmap)) / 8)) //nolint:gomnd
725 b := make([]byte, bitmapBytesLen)
726 for i := 0; i < len(bitmap); i++ {
727 if bitmap[i] {
728 b[i/8] |= 1 << (i % 8)
729 }
730 }
731 return b
732}
733
734func bytesToBitmap(b []byte) []bool {
735 var bitmap []bool
736 for i := 0; i < len(b); i++ {
737 for j := 0; j < 8; j++ {
738 bitmap = append(bitmap, b[i]&(1<<j) > 0)
739 }
740 }
741 return bitmap
742}
743
744// Get returns the value for a given key. If the key is not found, will return
745// the error ErrKeyNotFound, and in the leafK & leafV parameters will be placed
746// the data found in the tree in the leaf that was on the path going to the
747// input key.
748func (t *Tree) Get(k []byte) ([]byte, []byte, error) {
749 rTx := t.db.ReadTx()
750 defer rTx.Discard()
751
752 return t.GetWithTx(rTx, k)
753}
754
755// GetWithTx does the same than the Get method, but allowing to pass the
756// db.ReadTx that is used. If the key is not found, will return the error
757// ErrKeyNotFound, and in the leafK & leafV parameters will be placed the data
758// found in the tree in the leaf that was on the path going to the input key.
759func (t *Tree) GetWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, error) {
760 keyPath := make([]byte, t.hashFunction.Len())
761 copy(keyPath[:], k)
762
763 root, err := t.RootWithTx(rTx)
764 if err != nil {
765 return nil, nil, err
766 }
767
768 path := getPath(t.maxLevels, keyPath)
769 // go down to the leaf
770 var siblings [][]byte
771 _, value, _, err := t.down(rTx, k, root, siblings, path, 0, true)
772 if err != nil {
773 return nil, nil, err
774 }
775 leafK, leafV := ReadLeafValue(value)
776 if !bytes.Equal(k, leafK) {
777 return leafK, leafV, ErrKeyNotFound
778 }
779
780 return leafK, leafV, nil
781}
782
783// CheckProof verifies the given proof. The proof verification depends on the
784// HashFunction passed as parameter.
785func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool, error) {
786 siblings, err := UnpackSiblings(hashFunc, packedSiblings)
787 if err != nil {
788 return false, err
789 }
790
791 keyPath := make([]byte, hashFunc.Len())
792 copy(keyPath[:], k)
793
794 key, _, err := newLeafValue(hashFunc, k, v)
795 if err != nil {
796 return false, err
797 }
798
799 path := getPath(len(siblings), keyPath)
800 for i := len(siblings) - 1; i >= 0; i-- {
801 if path[i] {
802 key, _, err = newIntermediate(hashFunc, siblings[i], key)
803 if err != nil {
804 return false, err
805 }
806 } else {
807 key, _, err = newIntermediate(hashFunc, key, siblings[i])
808 if err != nil {
809 return false, err
810 }
811 }
812 }
813 if bytes.Equal(key[:], root) {
814 return true, nil
815 }
816 return false, nil
817}
818
819func (t *Tree) incNLeafs(wTx db.WriteTx, nLeafs int) error {
820 oldNLeafs, err := t.GetNLeafsWithTx(wTx)
821 if err != nil {
822 return err
823 }
824 newNLeafs := oldNLeafs + nLeafs
825 return t.setNLeafs(wTx, newNLeafs)
826}
827
828func (t *Tree) setNLeafs(wTx db.WriteTx, nLeafs int) error {
829 b := make([]byte, 8)
830 binary.LittleEndian.PutUint64(b, uint64(nLeafs))
831 if err := wTx.Set(dbKeyNLeafs, b); err != nil {
832 return err
833 }
834 return nil
835}
836
837// GetNLeafs returns the number of Leafs of the Tree.
838func (t *Tree) GetNLeafs() (int, error) {
839 rTx := t.db.ReadTx()
840 defer rTx.Discard()
841
842 return t.GetNLeafsWithTx(rTx)
843}
844
845// GetNLeafsWithTx does the same than the GetNLeafs method, but allowing to
846// pass the db.ReadTx that is used.
847func (t *Tree) GetNLeafsWithTx(rTx db.ReadTx) (int, error) {
848 b, err := rTx.Get(dbKeyNLeafs)
849 if err != nil {
850 return 0, err
851 }
852 nLeafs := binary.LittleEndian.Uint64(b)
853 return int(nLeafs), nil
854}
855
856// Snapshot returns a read-only copy of the Tree from the given root
857func (t *Tree) Snapshot(fromRoot []byte) (*Tree, error) {
858 // allow to define which root to use
859 if fromRoot == nil {
860 var err error
861 fromRoot, err = t.Root()
862 if err != nil {
863 return nil, err
864 }
865 }
866 return &Tree{
867 db: t.db,
868 maxLevels: t.maxLevels,
869 snapshotRoot: fromRoot,
870 emptyHash: t.emptyHash,
871 hashFunction: t.hashFunction,
872 dbg: t.dbg,
873 }, nil
874}
875
876// Iterate iterates through the full Tree, executing the given function on each
877// node of the Tree.
878func (t *Tree) Iterate(fromRoot []byte, f func([]byte, []byte)) error {
879 rTx := t.db.ReadTx()
880 defer rTx.Discard()
881
882 return t.IterateWithTx(rTx, fromRoot, f)
883}
884
885// IterateWithTx does the same than the Iterate method, but allowing to pass
886// the db.ReadTx that is used.
887func (t *Tree) IterateWithTx(rTx db.ReadTx, fromRoot []byte, f func([]byte, []byte)) error {
888 // allow to define which root to use
889 if fromRoot == nil {
890 var err error
891 fromRoot, err = t.RootWithTx(rTx)
892 if err != nil {
893 return err
894 }
895 }
896 return t.iter(rTx, fromRoot, f)
897}
898
899// IterateWithStop does the same than Iterate, but with int for the current
900// level, and a boolean parameter used by the passed function, is to indicate to
901// stop iterating on the branch when the method returns 'true'.
902func (t *Tree) IterateWithStop(fromRoot []byte, f func(int, []byte, []byte) bool) error {
903 rTx := t.db.ReadTx()
904 defer rTx.Discard()
905
906 // allow to define which root to use
907 if fromRoot == nil {
908 var err error
909 fromRoot, err = t.RootWithTx(rTx)
910 if err != nil {
911 return err
912 }
913 }
914 return t.iterWithStop(rTx, fromRoot, 0, f)
915}
916
917// IterateWithStopWithTx does the same than the IterateWithStop method, but
918// allowing to pass the db.ReadTx that is used.
919func (t *Tree) IterateWithStopWithTx(rTx db.ReadTx, fromRoot []byte,
920 f func(int, []byte, []byte) bool) error {
921 // allow to define which root to use
922 if fromRoot == nil {
923 var err error
924 fromRoot, err = t.RootWithTx(rTx)
925 if err != nil {
926 return err
927 }
928 }
929 return t.iterWithStop(rTx, fromRoot, 0, f)
930}
931
932func (t *Tree) iterWithStop(rTx db.ReadTx, k []byte, currLevel int,
933 f func(int, []byte, []byte) bool) error {
934 var v []byte
935 var err error
936 if bytes.Equal(k, t.emptyHash) {
937 v = t.emptyHash
938 } else {
939 v, err = rTx.Get(k)
940 if err != nil {
941 return err
942 }
943 }
944 currLevel++
945
946 switch v[0] {
947 case PrefixValueEmpty:
948 f(currLevel, k, v)
949 case PrefixValueLeaf:
950 f(currLevel, k, v)
951 case PrefixValueIntermediate:
952 stop := f(currLevel, k, v)
953 if stop {
954 return nil
955 }
956 l, r := ReadIntermediateChilds(v)
957 if err = t.iterWithStop(rTx, l, currLevel, f); err != nil {
958 return err
959 }
960 if err = t.iterWithStop(rTx, r, currLevel, f); err != nil {
961 return err
962 }
963 default:
964 return ErrInvalidValuePrefix
965 }
966 return nil
967}
968
969func (t *Tree) iter(rTx db.ReadTx, k []byte, f func([]byte, []byte)) error {
970 f2 := func(currLvl int, k, v []byte) bool {
971 f(k, v)
972 return false
973 }
974 return t.iterWithStop(rTx, k, 0, f2)
975}
976
977// Dump exports all the Tree leafs in a byte array of length:
978// [ N * (2+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
979// [ 1 byte | 1 byte | S bytes | len(v) bytes ]
980// [ len(k) | len(v) | key | value ]
981// Where S is the size of the output of the hash function used for the Tree.
982func (t *Tree) Dump(fromRoot []byte) ([]byte, error) {
983 // allow to define which root to use
984 if fromRoot == nil {
985 var err error
986 fromRoot, err = t.Root()
987 if err != nil {
988 return nil, err
989 }
990 }
991
992 // WARNING current encoding only supports key & values of 255 bytes each
993 // (due using only 1 byte for the length headers).
994 var b []byte
995 err := t.Iterate(fromRoot, func(k, v []byte) {
996 if v[0] != PrefixValueLeaf {
997 return
998 }
999 leafK, leafV := ReadLeafValue(v)
1000 kv := make([]byte, 2+len(leafK)+len(leafV))
1001 kv[0] = byte(len(leafK))
1002 kv[1] = byte(len(leafV))
1003 copy(kv[2:2+len(leafK)], leafK)
1004 copy(kv[2+len(leafK):], leafV)
1005 b = append(b, kv...)
1006 })
1007 return b, err
1008}
1009
1010// ImportDump imports the leafs (that have been exported with the Dump method)
1011// in the Tree.
1012func (t *Tree) ImportDump(b []byte) error {
1013 if !t.editable() {
1014 return ErrSnapshotNotEditable
1015 }
1016 root, err := t.Root()
1017 if err != nil {
1018 return err
1019 }
1020 if !bytes.Equal(root, t.emptyHash) {
1021 return ErrTreeNotEmpty
1022 }
1023
1024 r := bytes.NewReader(b)
1025 var keys, values [][]byte
1026 for {
1027 l := make([]byte, 2)
1028 _, err = io.ReadFull(r, l)
1029 if err == io.EOF {
1030 break
1031 } else if err != nil {
1032 return err
1033 }
1034 k := make([]byte, l[0])
1035 _, err = io.ReadFull(r, k)
1036 if err != nil {
1037 return err
1038 }
1039 v := make([]byte, l[1])
1040 _, err = io.ReadFull(r, v)
1041 if err != nil {
1042 return err
1043 }
1044 keys = append(keys, k)
1045 values = append(values, v)
1046 }
1047 if _, err = t.AddBatch(keys, values); err != nil {
1048 return err
1049 }
1050 return nil
1051}
1052
1053// Graphviz iterates across the full tree to generate a string Graphviz
1054// representation of the tree and writes it to w
1055func (t *Tree) Graphviz(w io.Writer, fromRoot []byte) error {
1056 return t.GraphvizFirstNLevels(w, fromRoot, t.maxLevels)
1057}
1058
1059// GraphvizFirstNLevels iterates across the first NLevels of the tree to
1060// generate a string Graphviz representation of the first NLevels of the tree
1061// and writes it to w
1062func (t *Tree) GraphvizFirstNLevels(w io.Writer, fromRoot []byte, untilLvl int) error {
1063 fmt.Fprintf(w, `digraph hierarchy {
1064node [fontname=Monospace,fontsize=10,shape=box]
1065`)
1066
1067 rTx := t.db.ReadTx()
1068 defer rTx.Discard()
1069
1070 if fromRoot == nil {
1071 var err error
1072 fromRoot, err = t.RootWithTx(rTx)
1073 if err != nil {
1074 return err
1075 }
1076 }
1077
1078 nEmpties := 0
1079 err := t.iterWithStop(rTx, fromRoot, 0, func(currLvl int, k, v []byte) bool {
1080 if currLvl == untilLvl {
1081 return true // to stop the iter from going down
1082 }
1083 switch v[0] {
1084 case PrefixValueEmpty:
1085 case PrefixValueLeaf:
1086 fmt.Fprintf(w, "\"%v\" [style=filled];\n", hex.EncodeToString(k[:nChars]))
1087 // key & value from the leaf
1088 kB, vB := ReadLeafValue(v)
1089 fmt.Fprintf(w, "\"%v\" -> {\"k:%v\\nv:%v\"}\n",
1090 hex.EncodeToString(k[:nChars]), hex.EncodeToString(kB[:nChars]),
1091 hex.EncodeToString(vB[:nChars]))
1092 fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
1093 hex.EncodeToString(kB[:nChars]), hex.EncodeToString(vB[:nChars]))
1094 case PrefixValueIntermediate:
1095 l, r := ReadIntermediateChilds(v)
1096 lStr := hex.EncodeToString(l[:nChars])
1097 rStr := hex.EncodeToString(r[:nChars])
1098 eStr := ""
1099 if bytes.Equal(l, t.emptyHash) {
1100 lStr = fmt.Sprintf("empty%v", nEmpties)
1101 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
1102 lStr)
1103 nEmpties++
1104 }
1105 if bytes.Equal(r, t.emptyHash) {
1106 rStr = fmt.Sprintf("empty%v", nEmpties)
1107 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
1108 rStr)
1109 nEmpties++
1110 }
1111 fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", hex.EncodeToString(k[:nChars]),
1112 lStr, rStr)
1113 fmt.Fprint(w, eStr)
1114 default:
1115 }
1116 return false
1117 })
1118 fmt.Fprintf(w, "}\n")
1119 return err
1120}
1121
1122// PrintGraphviz prints the output of Tree.Graphviz
1123func (t *Tree) PrintGraphviz(fromRoot []byte) error {
1124 if fromRoot == nil {
1125 var err error
1126 fromRoot, err = t.Root()
1127 if err != nil {
1128 return err
1129 }
1130 }
1131 return t.PrintGraphvizFirstNLevels(fromRoot, t.maxLevels)
1132}
1133
1134// PrintGraphvizFirstNLevels prints the output of Tree.GraphvizFirstNLevels
1135func (t *Tree) PrintGraphvizFirstNLevels(fromRoot []byte, untilLvl int) error {
1136 if fromRoot == nil {
1137 var err error
1138 fromRoot, err = t.Root()
1139 if err != nil {
1140 return err
1141 }
1142 }
1143 w := bytes.NewBufferString("")
1144 fmt.Fprintf(w,
1145 "--------\nGraphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+":\n")
1146 err := t.GraphvizFirstNLevels(w, fromRoot, untilLvl)
1147 if err != nil {
1148 fmt.Println(w)
1149 return err
1150 }
1151 fmt.Fprintf(w,
1152 "End of Graphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+"\n--------\n")
1153
1154 fmt.Println(w)
1155 return nil
1156}
1157
1158// TODO circom proofs
1159// TODO data structure for proofs (including root, key, value, siblings,
1160// hashFunction) + method to verify that data structure