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