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