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