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