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