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