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