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