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