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 "bufio"
15 "bytes"
16 "encoding/binary"
17 "encoding/hex"
18 "fmt"
19 "io"
20 "math"
21 "runtime"
22 "sync"
23
24 "go.vocdoni.io/dvote/db"
25)
26
27const (
28 // PrefixValueLen defines the bytes-prefix length used for the Value
29 // bytes representation stored in the db
30 PrefixValueLen = 2
31
32 // PrefixValueEmpty is used for the first byte of a Value to indicate
33 // that is an Empty value
34 PrefixValueEmpty = 0
35 // PrefixValueLeaf is used for the first byte of a Value to indicate
36 // that is a Leaf value
37 PrefixValueLeaf = 1
38 // PrefixValueIntermediate is used for the first byte of a Value to
39 // indicate that is a Intermediate value
40 PrefixValueIntermediate = 2
41
42 // nChars is used to crop the Graphviz nodes labels
43 nChars = 4
44
45 maxUint8 = int(^uint8(0)) // 2**8 -1
46 maxUint16 = int(^uint16(0)) // 2**16 -1
47)
48
49var (
50 // DefaultThresholdNLeafs defines the threshold number of leafs in the
51 // tree that determines if AddBatch will work in memory or in disk. It
52 // is defined when calling NewTree, and if set to 0 it will work always
53 // in disk.
54 DefaultThresholdNLeafs = 65536
55
56 dbKeyRoot = []byte("root")
57 dbKeyNLeafs = []byte("nleafs")
58 emptyValue = []byte{0}
59
60 // ErrKeyNotFound is used when a key is not found in the db neither in
61 // the current db Batch.
62 ErrKeyNotFound = fmt.Errorf("key not found")
63 // ErrKeyAlreadyExists is used when trying to add a key as leaf to the
64 // tree that already exists.
65 ErrKeyAlreadyExists = fmt.Errorf("key already exists")
66 // ErrInvalidValuePrefix is used when going down into the tree, a value
67 // is read from the db and has an unrecognized prefix.
68 ErrInvalidValuePrefix = fmt.Errorf("invalid value prefix")
69 // ErrDBNoTx is used when trying to use Tree.dbPut but Tree.dbBatch==nil
70 ErrDBNoTx = fmt.Errorf("dbPut error: no db Batch")
71 // ErrMaxLevel indicates when going down into the tree, the max level is
72 // reached
73 ErrMaxLevel = fmt.Errorf("max level reached")
74 // ErrMaxVirtualLevel indicates when going down into the tree, the max
75 // virtual level is reached
76 ErrMaxVirtualLevel = fmt.Errorf("max virtual level reached")
77 // ErrSnapshotNotEditable indicates when the tree is a non writable
78 // snapshot, thus can not be modified
79 ErrSnapshotNotEditable = fmt.Errorf("snapshot tree can not be edited")
80 // ErrTreeNotEmpty indicates when the tree was expected to be empty and
81 // it is not
82 ErrTreeNotEmpty = fmt.Errorf("tree is not empty")
83)
84
85// Tree defines the struct that implements the MerkleTree functionalities
86type Tree struct {
87 sync.Mutex
88
89 db db.Database
90 maxLevels int
91 // thresholdNLeafs defines the threshold number of leafs in the tree
92 // that determines if AddBatch will work in memory or in disk. It is
93 // defined when calling NewTree, and if set to 0 it will work always in
94 // disk.
95 thresholdNLeafs int
96 snapshotRoot []byte
97
98 hashFunction HashFunction
99 // TODO in the methods that use it, check if emptyHash param is len>0
100 // (check if it has been initialized)
101 emptyHash []byte
102
103 dbg *dbgStats
104}
105
106// Config defines the configuration for calling NewTree & NewTreeWithTx methods
107type Config struct {
108 Database db.Database
109 MaxLevels int
110 ThresholdNLeafs int
111 HashFunction HashFunction
112}
113
114// NewTree returns a new Tree, if there is a Tree still in the given database, it
115// will load it.
116func NewTree(cfg Config) (*Tree, error) {
117 wTx := cfg.Database.WriteTx()
118 defer wTx.Discard()
119
120 t, err := NewTreeWithTx(wTx, cfg)
121 if err != nil {
122 return nil, err
123 }
124
125 if err = wTx.Commit(); err != nil {
126 return nil, err
127 }
128 return t, nil
129}
130
131// NewTreeWithTx returns a new Tree using the given db.WriteTx, which will not
132// be ccommited inside this method, if there is a Tree still in the given
133// database, it will load it.
134func NewTreeWithTx(wTx db.WriteTx, cfg Config) (*Tree, error) {
135 // if thresholdNLeafs is set to 0, use the DefaultThresholdNLeafs
136 if cfg.ThresholdNLeafs == 0 {
137 cfg.ThresholdNLeafs = DefaultThresholdNLeafs
138 }
139 t := Tree{db: cfg.Database, maxLevels: cfg.MaxLevels,
140 thresholdNLeafs: cfg.ThresholdNLeafs, hashFunction: cfg.HashFunction}
141 t.emptyHash = make([]byte, t.hashFunction.Len()) // empty
142
143 _, err := wTx.Get(dbKeyRoot)
144 if err == db.ErrKeyNotFound {
145 // store new root 0 (empty)
146 if err = wTx.Set(dbKeyRoot, t.emptyHash); err != nil {
147 return nil, err
148 }
149 if err = t.setNLeafs(wTx, 0); err != nil {
150 return nil, err
151 }
152 return &t, nil
153 } else if err != nil {
154 return nil, err
155 }
156 return &t, nil
157}
158
159// Root returns the root of the Tree
160func (t *Tree) Root() ([]byte, error) {
161 rTx := t.db.ReadTx()
162 defer rTx.Discard()
163 return t.RootWithTx(rTx)
164}
165
166// RootWithTx returns the root of the Tree using the given db.ReadTx
167func (t *Tree) RootWithTx(rTx db.ReadTx) ([]byte, error) {
168 // if snapshotRoot is defined, means that the tree is a snapshot, and
169 // the root is not obtained from the db, but from the snapshotRoot
170 // parameter
171 if t.snapshotRoot != nil {
172 return t.snapshotRoot, nil
173 }
174 // get db root
175 return rTx.Get(dbKeyRoot)
176}
177
178func (t *Tree) setRoot(wTx db.WriteTx, root []byte) error {
179 return wTx.Set(dbKeyRoot, root)
180}
181
182// HashFunction returns Tree.hashFunction
183func (t *Tree) HashFunction() HashFunction {
184 return t.hashFunction
185}
186
187// editable returns true if the tree is editable, and false when is not
188// editable (because is a snapshot tree)
189func (t *Tree) editable() bool {
190 return t.snapshotRoot == nil
191}
192
193// Invalid is used when a key-value can not be added trough AddBatch, and
194// contains the index of the key-value and the error.
195type Invalid struct {
196 Index int
197 Error error
198}
199
200// AddBatch adds a batch of key-values to the Tree. Returns an array containing
201// the indexes of the keys failed to add. Supports empty values as input
202// parameters, which is equivalent to 0 valued byte array.
203func (t *Tree) AddBatch(keys, values [][]byte) ([]Invalid, error) {
204 wTx := t.db.WriteTx()
205 defer wTx.Discard()
206
207 invalids, err := t.AddBatchWithTx(wTx, keys, values)
208 if err != nil {
209 return invalids, err
210 }
211 return invalids, wTx.Commit()
212}
213
214// AddBatchWithTx does the same than the AddBatch method, but allowing to pass
215// the db.WriteTx that is used. The db.WriteTx will not be committed inside
216// this method.
217func (t *Tree) AddBatchWithTx(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
218 t.Lock()
219 defer t.Unlock()
220
221 if !t.editable() {
222 return nil, ErrSnapshotNotEditable
223 }
224
225 e := []byte{}
226 // equal the number of keys & values
227 if len(keys) > len(values) {
228 // add missing values
229 for i := len(values); i < len(keys); i++ {
230 values = append(values, e)
231 }
232 } else if len(keys) < len(values) {
233 // crop extra values
234 values = values[:len(keys)]
235 }
236
237 nLeafs, err := t.GetNLeafsWithTx(wTx)
238 if err != nil {
239 return nil, err
240 }
241 if nLeafs > t.thresholdNLeafs {
242 return t.addBatchInDisk(wTx, keys, values)
243 }
244 return t.addBatchInMemory(wTx, keys, values)
245}
246
247func (t *Tree) addBatchInDisk(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
248 nCPU := flp2(runtime.NumCPU())
249 if nCPU == 1 || len(keys) < nCPU {
250 var invalids []Invalid
251 for i := 0; i < len(keys); i++ {
252 if err := t.addWithTx(wTx, keys[i], values[i]); err != nil {
253 invalids = append(invalids, Invalid{i, err})
254 }
255 }
256 return invalids, nil
257 }
258
259 kvs, invalids, err := keysValuesToKvs(t.maxLevels, keys, values)
260 if err != nil {
261 return nil, err
262 }
263
264 buckets := splitInBuckets(kvs, nCPU)
265
266 root, err := t.RootWithTx(wTx)
267 if err != nil {
268 return nil, err
269 }
270
271 l := int(math.Log2(float64(nCPU)))
272 subRoots, err := t.getSubRootsAtLevel(wTx, root, l+1)
273 if err != nil {
274 return nil, err
275 }
276 if len(subRoots) != nCPU {
277 // Already populated Tree but Unbalanced.
278
279 // add one key at each bucket, and then continue with the flow
280 for i := 0; i < len(buckets); i++ {
281 // add one leaf of the bucket, if there is an error when
282 // adding the k-v, try to add the next one of the bucket
283 // (until one is added)
284 inserted := -1
285 for j := 0; j < len(buckets[i]); j++ {
286 if newRoot, err := t.add(wTx, root, 0,
287 buckets[i][j].k, buckets[i][j].v); err == nil {
288 inserted = j
289 root = newRoot
290 break
291 }
292 }
293
294 // remove the inserted element from buckets[i]
295 if inserted != -1 {
296 buckets[i] = append(buckets[i][:inserted], buckets[i][inserted+1:]...)
297 }
298 }
299 subRoots, err = t.getSubRootsAtLevel(wTx, root, l+1)
300 if err != nil {
301 return nil, err
302 }
303 }
304
305 if len(subRoots) != nCPU {
306 return nil, fmt.Errorf("This error should not be reached."+
307 " len(subRoots) != nCPU, len(subRoots)=%d, nCPU=%d."+
308 " Please report it in a new issue:"+
309 " https://github.com/vocdoni/arbo/issues/new", len(subRoots), nCPU)
310 }
311
312 invalidsInBucket := make([][]Invalid, nCPU)
313 txs := make([]db.WriteTx, nCPU)
314 for i := 0; i < nCPU; i++ {
315 txs[i] = t.db.WriteTx()
316 err := txs[i].Apply(wTx)
317 if err != nil {
318 return nil, err
319 }
320 }
321
322 var wg sync.WaitGroup
323 wg.Add(nCPU)
324 for i := 0; i < nCPU; i++ {
325 go func(cpu int) {
326 // use different wTx for each cpu, after once all
327 // are done, iter over the cpuWTxs and copy their
328 // content into the main wTx
329 for j := 0; j < len(buckets[cpu]); j++ {
330 newSubRoot, err := t.add(txs[cpu], subRoots[cpu],
331 l, buckets[cpu][j].k, buckets[cpu][j].v)
332 if err != nil {
333 invalidsInBucket[cpu] = append(invalidsInBucket[cpu],
334 Invalid{buckets[cpu][j].pos, err})
335 continue
336 }
337 // if there has not been errors, set the new subRoots[cpu]
338 subRoots[cpu] = newSubRoot
339 }
340 wg.Done()
341 }(i)
342 }
343 wg.Wait()
344
345 for i := 0; i < nCPU; i++ {
346 if err := wTx.Apply(txs[i]); err != nil {
347 return nil, err
348 }
349 txs[i].Discard()
350 }
351
352 for i := 0; i < len(invalidsInBucket); i++ {
353 invalids = append(invalids, invalidsInBucket[i]...)
354 }
355
356 newRoot, err := t.upFromSubRoots(wTx, subRoots)
357 if err != nil {
358 return nil, err
359 }
360
361 // update dbKeyNLeafs
362 if err := t.SetRootWithTx(wTx, newRoot); err != nil {
363 return nil, err
364 }
365
366 // update nLeafs
367 if err := t.incNLeafs(wTx, len(keys)-len(invalids)); err != nil {
368 return nil, err
369 }
370
371 return invalids, nil
372}
373
374func (t *Tree) upFromSubRoots(wTx db.WriteTx, subRoots [][]byte) ([]byte, error) {
375 // is a method of Tree just to get access to t.hashFunction and
376 // t.emptyHash.
377
378 // go up from subRoots to up, storing nodes in the given WriteTx
379 // once up at the root, store it in the WriteTx using the dbKeyRoot
380 if len(subRoots) == 1 {
381 return subRoots[0], nil
382 }
383 // get the subRoots values to know the node types of each subRoot
384 nodeTypes := make([]byte, len(subRoots))
385 for i := 0; i < len(subRoots); i++ {
386 if bytes.Equal(subRoots[i], t.emptyHash) {
387 nodeTypes[i] = PrefixValueEmpty
388 continue
389 }
390 v, err := wTx.Get(subRoots[i])
391 if err != nil {
392 return nil, err
393 }
394 nodeTypes[i] = v[0]
395 }
396
397 var newSubRoots [][]byte
398 for i := 0; i < len(subRoots); i += 2 {
399 if (bytes.Equal(subRoots[i], t.emptyHash) && bytes.Equal(subRoots[i+1], t.emptyHash)) ||
400 (nodeTypes[i] == PrefixValueLeaf && bytes.Equal(subRoots[i+1], t.emptyHash)) {
401 // when both sub nodes are empty, the parent is also empty
402 // or
403 // when 1st sub node is a leaf but the 2nd is empty, the
404 // leaf is used as 'parent'
405
406 newSubRoots = append(newSubRoots, subRoots[i])
407 continue
408 }
409 if bytes.Equal(subRoots[i], t.emptyHash) && nodeTypes[i+1] == PrefixValueLeaf {
410 // when 2nd sub node is a leaf but the 1st is empty,
411 // the leaf is used as 'parent'
412 newSubRoots = append(newSubRoots, subRoots[i+1])
413 continue
414 }
415
416 k, v, err := t.newIntermediate(subRoots[i], subRoots[i+1])
417 if err != nil {
418 return nil, err
419 }
420 // store k-v to db
421 if err = wTx.Set(k, v); err != nil {
422 return nil, err
423 }
424 newSubRoots = append(newSubRoots, k)
425 }
426
427 return t.upFromSubRoots(wTx, newSubRoots)
428}
429
430func (t *Tree) getSubRootsAtLevel(rTx db.ReadTx, root []byte, l int) ([][]byte, error) {
431 // go at level l and return each node key, where each node key is the
432 // subRoot of the subTree that starts there
433
434 var subRoots [][]byte
435 err := t.iterWithStop(rTx, root, 0, func(currLvl int, k, v []byte) bool {
436 if currLvl == l && !bytes.Equal(k, t.emptyHash) {
437 subRoots = append(subRoots, k)
438 }
439 if currLvl >= l {
440 return true // to stop the iter from going down
441 }
442 return false
443 })
444
445 return subRoots, err
446}
447
448func (t *Tree) addBatchInMemory(wTx db.WriteTx, keys, values [][]byte) ([]Invalid, error) {
449 vt, err := t.loadVT(wTx)
450 if err != nil {
451 return nil, err
452 }
453
454 invalids, err := vt.addBatch(keys, values)
455 if err != nil {
456 return nil, err
457 }
458
459 // once the VirtualTree is build, compute the hashes
460 pairs, err := vt.computeHashes()
461 if err != nil {
462 // currently invalids in computeHashes are not counted,
463 // but should not be needed, as if there is an error there is
464 // nothing stored in the db and the error is returned
465 return nil, err
466 }
467
468 // store pairs in db
469 for i := 0; i < len(pairs); i++ {
470 if err := wTx.Set(pairs[i][0], pairs[i][1]); err != nil {
471 return nil, err
472 }
473 }
474
475 // store root (from the vt) to db
476 if vt.root != nil {
477 if err := wTx.Set(dbKeyRoot, vt.root.h); err != nil {
478 return nil, err
479 }
480 }
481
482 // update nLeafs
483 if err := t.incNLeafs(wTx, len(keys)-len(invalids)); err != nil {
484 return nil, err
485 }
486
487 return invalids, nil
488}
489
490// loadVT loads a new virtual tree (vt) from the current Tree, which contains
491// the same leafs.
492func (t *Tree) loadVT(rTx db.ReadTx) (vt, error) {
493 vt := newVT(t.maxLevels, t.hashFunction)
494 vt.params.dbg = t.dbg
495 var callbackErr error
496 err := t.IterateWithStopWithTx(rTx, nil, func(_ int, k, v []byte) bool {
497 if v[0] != PrefixValueLeaf {
498 return false
499 }
500 leafK, leafV := ReadLeafValue(v)
501 if err := vt.add(0, leafK, leafV); err != nil {
502 callbackErr = err
503 return true
504 }
505 return false
506 })
507 if callbackErr != nil {
508 return vt, callbackErr
509 }
510
511 return vt, err
512}
513
514// Add inserts the key-value into the Tree. If the inputs come from a
515// *big.Int, is expected that are represented by a Little-Endian byte array
516// (for circom compatibility).
517func (t *Tree) Add(k, v []byte) error {
518 wTx := t.db.WriteTx()
519 defer wTx.Discard()
520
521 if err := t.AddWithTx(wTx, k, v); err != nil {
522 return err
523 }
524
525 return wTx.Commit()
526}
527
528// AddWithTx does the same than the Add method, but allowing to pass the
529// db.WriteTx that is used. The db.WriteTx will not be committed inside this
530// method.
531func (t *Tree) AddWithTx(wTx db.WriteTx, k, v []byte) error {
532 t.Lock()
533 defer t.Unlock()
534
535 if !t.editable() {
536 return ErrSnapshotNotEditable
537 }
538 return t.addWithTx(wTx, k, v)
539}
540
541// warning: addWithTx does not use the Tree mutex, the mutex is responsibility
542// of the methods calling this method, and same with t.editable().
543func (t *Tree) addWithTx(wTx db.WriteTx, k, v []byte) error {
544 root, err := t.RootWithTx(wTx)
545 if err != nil {
546 return err
547 }
548
549 root, err = t.add(wTx, root, 0, k, v) // add from level 0
550 if err != nil {
551 return err
552 }
553 // store root to db
554 if err := t.setRoot(wTx, root); err != nil {
555 return err
556 }
557 // update nLeafs
558 if err = t.incNLeafs(wTx, 1); err != nil {
559 return err
560 }
561 return nil
562}
563
564// keyPathFromKey returns the keyPath and checks that the key is not bigger
565// than maximum key length for the tree maxLevels size.
566// This is because if the key bits length is bigger than the maxLevels of the
567// tree, two different keys that their difference is at the end, will collision
568// in the same leaf of the tree (at the max depth).
569func keyPathFromKey(maxLevels int, k []byte) ([]byte, error) {
570 maxKeyLen := int(math.Ceil(float64(maxLevels) / float64(8))) //nolint:gomnd
571 if len(k) > maxKeyLen {
572 return nil, fmt.Errorf("len(k) can not be bigger than ceil(maxLevels/8), where"+
573 " len(k): %d, maxLevels: %d, max key len=ceil(maxLevels/8): %d. Might need"+
574 " a bigger tree depth (maxLevels>=%d) in order to input keys of length %d",
575 len(k), maxLevels, maxKeyLen, len(k)*8, len(k)) //nolint:gomnd
576 }
577 keyPath := make([]byte, maxKeyLen) //nolint:gomnd
578 copy(keyPath[:], k)
579 return keyPath, nil
580}
581
582func (t *Tree) add(wTx db.WriteTx, root []byte, fromLvl int, k, v []byte) ([]byte, error) {
583 keyPath, err := keyPathFromKey(t.maxLevels, k)
584 if err != nil {
585 return nil, err
586 }
587 path := getPath(t.maxLevels, keyPath)
588
589 // go down to the leaf
590 var siblings [][]byte
591 _, _, siblings, err = t.down(wTx, k, root, siblings, path, fromLvl, false)
592 if err != nil {
593 return nil, err
594 }
595
596 leafKey, leafValue, err := t.newLeafValue(k, v)
597 if err != nil {
598 return nil, err
599 }
600
601 if err := wTx.Set(leafKey, leafValue); err != nil {
602 return nil, err
603 }
604
605 // go up to the root
606 if len(siblings) == 0 {
607 // return the leafKey as root
608 return leafKey, nil
609 }
610 root, err = t.up(wTx, leafKey, siblings, path, len(siblings)-1, fromLvl)
611 if err != nil {
612 return nil, err
613 }
614
615 return root, nil
616}
617
618// down goes down to the leaf recursively
619func (t *Tree) down(rTx db.ReadTx, newKey, currKey []byte, siblings [][]byte,
620 path []bool, currLvl int, getLeaf bool) (
621 []byte, []byte, [][]byte, error) {
622 if currLvl > t.maxLevels {
623 return nil, nil, nil, ErrMaxLevel
624 }
625
626 var err error
627 var currValue []byte
628 if bytes.Equal(currKey, t.emptyHash) {
629 // empty value
630 return currKey, emptyValue, siblings, nil
631 }
632 currValue, err = rTx.Get(currKey)
633 if err != nil {
634 return nil, nil, nil, err
635 }
636
637 switch currValue[0] {
638 case PrefixValueEmpty: // empty
639 fmt.Printf("newKey: %s, currKey: %s, currLvl: %d, currValue: %s\n",
640 hex.EncodeToString(newKey), hex.EncodeToString(currKey),
641 currLvl, hex.EncodeToString(currValue))
642 panic("This point should not be reached, as the 'if currKey==t.emptyHash'" +
643 " above should avoid reaching this point. This panic is temporary" +
644 " for reporting purposes, will be deleted in future versions." +
645 " Please paste this log (including the previous log lines) in a" +
646 " new issue: https://github.com/vocdoni/arbo/issues/new") // TMP
647 case PrefixValueLeaf: // leaf
648 if !bytes.Equal(currValue, emptyValue) {
649 if getLeaf {
650 return currKey, currValue, siblings, nil
651 }
652 oldLeafKey, _ := ReadLeafValue(currValue)
653 if bytes.Equal(newKey, oldLeafKey) {
654 return nil, nil, nil, ErrKeyAlreadyExists
655 }
656
657 oldLeafKeyFull, err := keyPathFromKey(t.maxLevels, oldLeafKey)
658 if err != nil {
659 return nil, nil, nil, err
660 }
661
662 // if currKey is already used, go down until paths diverge
663 oldPath := getPath(t.maxLevels, oldLeafKeyFull)
664 siblings, err = t.downVirtually(siblings, currKey, newKey, oldPath, path, currLvl)
665 if err != nil {
666 return nil, nil, nil, err
667 }
668 }
669 return currKey, currValue, siblings, nil
670 case PrefixValueIntermediate: // intermediate
671 if len(currValue) != PrefixValueLen+t.hashFunction.Len()*2 {
672 return nil, nil, nil,
673 fmt.Errorf("intermediate value invalid length (expected: %d, actual: %d)",
674 PrefixValueLen+t.hashFunction.Len()*2, len(currValue))
675 }
676 // collect siblings while going down
677 if path[currLvl] {
678 // right
679 lChild, rChild := ReadIntermediateChilds(currValue)
680 siblings = append(siblings, lChild)
681 return t.down(rTx, newKey, rChild, siblings, path, currLvl+1, getLeaf)
682 }
683 // left
684 lChild, rChild := ReadIntermediateChilds(currValue)
685 siblings = append(siblings, rChild)
686 return t.down(rTx, newKey, lChild, siblings, path, currLvl+1, getLeaf)
687 default:
688 return nil, nil, nil, ErrInvalidValuePrefix
689 }
690}
691
692// downVirtually is used when in a leaf already exists, and a new leaf which
693// shares the path until the existing leaf is being added
694func (t *Tree) downVirtually(siblings [][]byte, oldKey, newKey []byte, oldPath,
695 newPath []bool, currLvl int) ([][]byte, error) {
696 var err error
697 if currLvl > t.maxLevels-1 {
698 return nil, ErrMaxVirtualLevel
699 }
700
701 if oldPath[currLvl] == newPath[currLvl] {
702 siblings = append(siblings, t.emptyHash)
703
704 siblings, err = t.downVirtually(siblings, oldKey, newKey, oldPath, newPath, currLvl+1)
705 if err != nil {
706 return nil, err
707 }
708 return siblings, nil
709 }
710 // reached the divergence
711 siblings = append(siblings, oldKey)
712
713 return siblings, nil
714}
715
716// up goes up recursively updating the intermediate nodes
717func (t *Tree) up(wTx db.WriteTx, key []byte, siblings [][]byte, path []bool,
718 currLvl, toLvl int) ([]byte, error) {
719 var k, v []byte
720 var err error
721 if path[currLvl+toLvl] {
722 k, v, err = t.newIntermediate(siblings[currLvl], key)
723 if err != nil {
724 return nil, err
725 }
726 } else {
727 k, v, err = t.newIntermediate(key, siblings[currLvl])
728 if err != nil {
729 return nil, err
730 }
731 }
732 // store k-v to db
733 if err = wTx.Set(k, v); err != nil {
734 return nil, err
735 }
736
737 if currLvl == 0 {
738 // reached the root
739 return k, nil
740 }
741
742 return t.up(wTx, k, siblings, path, currLvl-1, toLvl)
743}
744
745func (t *Tree) newLeafValue(k, v []byte) ([]byte, []byte, error) {
746 t.dbg.incHash()
747 return newLeafValue(t.hashFunction, k, v)
748}
749
750// newLeafValue takes a key & value from a leaf, and computes the leaf hash,
751// which is used as the leaf key. And the value is the concatenation of the
752// inputed key & value. The output of this function is used as key-value to
753// store the leaf in the DB.
754// [ 1 byte | 1 byte | N bytes | M bytes ]
755// [ type of node | length of key | key | value ]
756func newLeafValue(hashFunc HashFunction, k, v []byte) ([]byte, []byte, error) {
757 leafKey, err := hashFunc.Hash(k, v, []byte{1})
758 if err != nil {
759 return nil, nil, err
760 }
761 var leafValue []byte
762 leafValue = append(leafValue, byte(PrefixValueLeaf))
763 if len(k) > maxUint8 {
764 return nil, nil, fmt.Errorf("newLeafValue: len(k) > %v", maxUint8)
765 }
766 leafValue = append(leafValue, byte(len(k)))
767 leafValue = append(leafValue, k...)
768 leafValue = append(leafValue, v...)
769 return leafKey, leafValue, nil
770}
771
772// ReadLeafValue reads from a byte array the leaf key & value
773func ReadLeafValue(b []byte) ([]byte, []byte) {
774 if len(b) < PrefixValueLen {
775 return []byte{}, []byte{}
776 }
777
778 kLen := b[1]
779 if len(b) < PrefixValueLen+int(kLen) {
780 return []byte{}, []byte{}
781 }
782 k := b[PrefixValueLen : PrefixValueLen+kLen]
783 v := b[PrefixValueLen+kLen:]
784 return k, v
785}
786
787func (t *Tree) newIntermediate(l, r []byte) ([]byte, []byte, error) {
788 t.dbg.incHash()
789 return newIntermediate(t.hashFunction, l, r)
790}
791
792// newIntermediate takes the left & right keys of a intermediate node, and
793// computes its hash. Returns the hash of the node, which is the node key, and a
794// byte array that contains the value (which contains the left & right child
795// keys) to store in the DB.
796// [ 1 byte | 1 byte | N bytes | N bytes ]
797// [ type of node | length of left key | left key | right key ]
798func newIntermediate(hashFunc HashFunction, l, r []byte) ([]byte, []byte, error) {
799 b := make([]byte, PrefixValueLen+hashFunc.Len()*2)
800 b[0] = PrefixValueIntermediate
801 if len(l) > maxUint8 {
802 return nil, nil, fmt.Errorf("newIntermediate: len(l) > %v", maxUint8)
803 }
804 b[1] = byte(len(l))
805 copy(b[PrefixValueLen:PrefixValueLen+hashFunc.Len()], l)
806 copy(b[PrefixValueLen+hashFunc.Len():], r)
807
808 key, err := hashFunc.Hash(l, r)
809 if err != nil {
810 return nil, nil, err
811 }
812
813 return key, b, nil
814}
815
816// ReadIntermediateChilds reads from a byte array the two childs keys
817func ReadIntermediateChilds(b []byte) ([]byte, []byte) {
818 if len(b) < PrefixValueLen {
819 return []byte{}, []byte{}
820 }
821
822 lLen := b[1]
823 if len(b) < PrefixValueLen+int(lLen) {
824 return []byte{}, []byte{}
825 }
826 l := b[PrefixValueLen : PrefixValueLen+lLen]
827 r := b[PrefixValueLen+lLen:]
828 return l, r
829}
830
831func getPath(numLevels int, k []byte) []bool {
832 path := make([]bool, numLevels)
833 for n := 0; n < numLevels; n++ {
834 path[n] = k[n/8]&(1<<(n%8)) != 0
835 }
836 return path
837}
838
839// Update updates the value for a given existing key. If the given key does not
840// exist, returns an error.
841func (t *Tree) Update(k, v []byte) error {
842 wTx := t.db.WriteTx()
843 defer wTx.Discard()
844
845 if err := t.UpdateWithTx(wTx, k, v); err != nil {
846 return err
847 }
848 return wTx.Commit()
849}
850
851// UpdateWithTx does the same than the Update method, but allowing to pass the
852// db.WriteTx that is used. The db.WriteTx will not be committed inside this
853// method.
854func (t *Tree) UpdateWithTx(wTx db.WriteTx, k, v []byte) error {
855 t.Lock()
856 defer t.Unlock()
857
858 if !t.editable() {
859 return ErrSnapshotNotEditable
860 }
861
862 keyPath, err := keyPathFromKey(t.maxLevels, k)
863 if err != nil {
864 return err
865 }
866 path := getPath(t.maxLevels, keyPath)
867
868 root, err := t.RootWithTx(wTx)
869 if err != nil {
870 return err
871 }
872
873 var siblings [][]byte
874 _, valueAtBottom, siblings, err := t.down(wTx, k, root, siblings, path, 0, true)
875 if err != nil {
876 return err
877 }
878 oldKey, _ := ReadLeafValue(valueAtBottom)
879 if !bytes.Equal(oldKey, k) {
880 return ErrKeyNotFound
881 }
882
883 leafKey, leafValue, err := t.newLeafValue(k, v)
884 if err != nil {
885 return err
886 }
887
888 if err := wTx.Set(leafKey, leafValue); err != nil {
889 return err
890 }
891
892 // go up to the root
893 if len(siblings) == 0 {
894 return t.setRoot(wTx, leafKey)
895 }
896 root, err = t.up(wTx, leafKey, siblings, path, len(siblings)-1, 0)
897 if err != nil {
898 return err
899 }
900
901 // store root to db
902 if err := t.setRoot(wTx, root); err != nil {
903 return err
904 }
905 return nil
906}
907
908// GenProof generates a MerkleTree proof for the given key. The leaf value is
909// returned, together with the packed siblings of the proof, and a boolean
910// parameter that indicates if the proof is of existence (true) or not (false).
911func (t *Tree) GenProof(k []byte) ([]byte, []byte, []byte, bool, error) {
912 rTx := t.db.ReadTx()
913 defer rTx.Discard()
914
915 return t.GenProofWithTx(rTx, k)
916}
917
918// GenProofWithTx does the same than the GenProof method, but allowing to pass
919// the db.ReadTx that is used.
920func (t *Tree) GenProofWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, []byte, bool, error) {
921 keyPath, err := keyPathFromKey(t.maxLevels, k)
922 if err != nil {
923 return nil, nil, nil, false, err
924 }
925 path := getPath(t.maxLevels, keyPath)
926
927 root, err := t.RootWithTx(rTx)
928 if err != nil {
929 return nil, nil, nil, false, err
930 }
931
932 // go down to the leaf
933 var siblings [][]byte
934 _, value, siblings, err := t.down(rTx, k, root, siblings, path, 0, true)
935 if err != nil {
936 return nil, nil, nil, false, err
937 }
938
939 s, err := PackSiblings(t.hashFunction, siblings)
940 if err != nil {
941 return nil, nil, nil, false, err
942 }
943
944 leafK, leafV := ReadLeafValue(value)
945 if !bytes.Equal(k, leafK) {
946 // key not in tree, proof of non-existence
947 return leafK, leafV, s, false, nil
948 }
949
950 return leafK, leafV, s, true, nil
951}
952
953// PackSiblings packs the siblings into a byte array.
954// [ 2 byte | 2 byte | L bytes | S * N bytes ]
955// [ full length | bitmap length (L) | bitmap | N non-zero siblings ]
956// Where the bitmap indicates if the sibling is 0 or a value from the siblings
957// array. And S is the size of the output of the hash function used for the
958// Tree. The 2 2-byte that define the full length and bitmap length, are
959// encoded in little-endian.
960func PackSiblings(hashFunc HashFunction, siblings [][]byte) ([]byte, error) {
961 var b []byte
962 var bitmap []bool
963 emptySibling := make([]byte, hashFunc.Len())
964 for i := 0; i < len(siblings); i++ {
965 if bytes.Equal(siblings[i], emptySibling) {
966 bitmap = append(bitmap, false)
967 } else {
968 bitmap = append(bitmap, true)
969 b = append(b, siblings[i]...)
970 }
971 }
972
973 bitmapBytes := bitmapToBytes(bitmap)
974 l := len(bitmapBytes)
975 if l > maxUint16 {
976 return nil, fmt.Errorf("PackSiblings: bitmapBytes length > %v", maxUint16)
977 }
978
979 fullLen := 4 + l + len(b) //nolint:gomnd
980 if fullLen > maxUint16 {
981 return nil, fmt.Errorf("PackSiblings: fullLen > %v", maxUint16)
982 }
983 res := make([]byte, fullLen)
984 binary.LittleEndian.PutUint16(res[0:2], uint16(fullLen)) // set full length
985 binary.LittleEndian.PutUint16(res[2:4], uint16(l)) // set the bitmapBytes length
986 copy(res[4:4+l], bitmapBytes)
987 copy(res[4+l:], b)
988 return res, nil
989}
990
991// UnpackSiblings unpacks the siblings from a byte array.
992func UnpackSiblings(hashFunc HashFunction, b []byte) ([][]byte, error) {
993 fullLen := binary.LittleEndian.Uint16(b[0:2])
994 l := binary.LittleEndian.Uint16(b[2:4]) // bitmap bytes length
995 if len(b) != int(fullLen) {
996 return nil,
997 fmt.Errorf("expected len: %d, current len: %d",
998 fullLen, len(b))
999 }
1000
1001 bitmapBytes := b[4 : 4+l]
1002 bitmap := bytesToBitmap(bitmapBytes)
1003 siblingsBytes := b[4+l:]
1004 iSibl := 0
1005 emptySibl := make([]byte, hashFunc.Len())
1006 var siblings [][]byte
1007 for i := 0; i < len(bitmap); i++ {
1008 if iSibl >= len(siblingsBytes) {
1009 break
1010 }
1011 if bitmap[i] {
1012 siblings = append(siblings, siblingsBytes[iSibl:iSibl+hashFunc.Len()])
1013 iSibl += hashFunc.Len()
1014 } else {
1015 siblings = append(siblings, emptySibl)
1016 }
1017 }
1018 return siblings, nil
1019}
1020
1021func bitmapToBytes(bitmap []bool) []byte {
1022 bitmapBytesLen := int(math.Ceil(float64(len(bitmap)) / 8)) //nolint:gomnd
1023 b := make([]byte, bitmapBytesLen)
1024 for i := 0; i < len(bitmap); i++ {
1025 if bitmap[i] {
1026 b[i/8] |= 1 << (i % 8)
1027 }
1028 }
1029 return b
1030}
1031
1032func bytesToBitmap(b []byte) []bool {
1033 var bitmap []bool
1034 for i := 0; i < len(b); i++ {
1035 for j := 0; j < 8; j++ {
1036 bitmap = append(bitmap, b[i]&(1<<j) > 0)
1037 }
1038 }
1039 return bitmap
1040}
1041
1042// Get returns the value in the Tree for a given key. If the key is not found,
1043// will return the error ErrKeyNotFound, and in the leafK & leafV parameters
1044// will be placed the data found in the tree in the leaf that was on the path
1045// going to the input key.
1046func (t *Tree) Get(k []byte) ([]byte, []byte, error) {
1047 rTx := t.db.ReadTx()
1048 defer rTx.Discard()
1049
1050 return t.GetWithTx(rTx, k)
1051}
1052
1053// GetWithTx does the same than the Get method, but allowing to pass the
1054// db.ReadTx that is used. If the key is not found, will return the error
1055// ErrKeyNotFound, and in the leafK & leafV parameters will be placed the data
1056// found in the tree in the leaf that was on the path going to the input key.
1057func (t *Tree) GetWithTx(rTx db.ReadTx, k []byte) ([]byte, []byte, error) {
1058 keyPath, err := keyPathFromKey(t.maxLevels, k)
1059 if err != nil {
1060 return nil, nil, err
1061 }
1062 path := getPath(t.maxLevels, keyPath)
1063
1064 root, err := t.RootWithTx(rTx)
1065 if err != nil {
1066 return nil, nil, err
1067 }
1068
1069 // go down to the leaf
1070 var siblings [][]byte
1071 _, value, _, err := t.down(rTx, k, root, siblings, path, 0, true)
1072 if err != nil {
1073 return nil, nil, err
1074 }
1075 leafK, leafV := ReadLeafValue(value)
1076 if !bytes.Equal(k, leafK) {
1077 return leafK, leafV, ErrKeyNotFound
1078 }
1079
1080 return leafK, leafV, nil
1081}
1082
1083// CheckProof verifies the given proof. The proof verification depends on the
1084// HashFunction passed as parameter.
1085func CheckProof(hashFunc HashFunction, k, v, root, packedSiblings []byte) (bool, error) {
1086 siblings, err := UnpackSiblings(hashFunc, packedSiblings)
1087 if err != nil {
1088 return false, err
1089 }
1090
1091 keyPath := make([]byte, int(math.Ceil(float64(len(siblings))/float64(8)))) //nolint:gomnd
1092 copy(keyPath[:], k)
1093
1094 key, _, err := newLeafValue(hashFunc, k, v)
1095 if err != nil {
1096 return false, err
1097 }
1098
1099 path := getPath(len(siblings), keyPath)
1100 for i := len(siblings) - 1; i >= 0; i-- {
1101 if path[i] {
1102 key, _, err = newIntermediate(hashFunc, siblings[i], key)
1103 if err != nil {
1104 return false, err
1105 }
1106 } else {
1107 key, _, err = newIntermediate(hashFunc, key, siblings[i])
1108 if err != nil {
1109 return false, err
1110 }
1111 }
1112 }
1113 if bytes.Equal(key[:], root) {
1114 return true, nil
1115 }
1116 return false, nil
1117}
1118
1119func (t *Tree) incNLeafs(wTx db.WriteTx, nLeafs int) error {
1120 oldNLeafs, err := t.GetNLeafsWithTx(wTx)
1121 if err != nil {
1122 return err
1123 }
1124 newNLeafs := oldNLeafs + nLeafs
1125 return t.setNLeafs(wTx, newNLeafs)
1126}
1127
1128func (t *Tree) setNLeafs(wTx db.WriteTx, nLeafs int) error {
1129 b := make([]byte, 8)
1130 binary.LittleEndian.PutUint64(b, uint64(nLeafs))
1131 if err := wTx.Set(dbKeyNLeafs, b); err != nil {
1132 return err
1133 }
1134 return nil
1135}
1136
1137// GetNLeafs returns the number of Leafs of the Tree.
1138func (t *Tree) GetNLeafs() (int, error) {
1139 rTx := t.db.ReadTx()
1140 defer rTx.Discard()
1141
1142 return t.GetNLeafsWithTx(rTx)
1143}
1144
1145// GetNLeafsWithTx does the same than the GetNLeafs method, but allowing to
1146// pass the db.ReadTx that is used.
1147func (t *Tree) GetNLeafsWithTx(rTx db.ReadTx) (int, error) {
1148 b, err := rTx.Get(dbKeyNLeafs)
1149 if err != nil {
1150 return 0, err
1151 }
1152 nLeafs := binary.LittleEndian.Uint64(b)
1153 return int(nLeafs), nil
1154}
1155
1156// SetRoot sets the root to the given root
1157func (t *Tree) SetRoot(root []byte) error {
1158 wTx := t.db.WriteTx()
1159 defer wTx.Discard()
1160
1161 if err := t.SetRootWithTx(wTx, root); err != nil {
1162 return err
1163 }
1164 return wTx.Commit()
1165}
1166
1167// SetRootWithTx sets the root to the given root using the given db.WriteTx
1168func (t *Tree) SetRootWithTx(wTx db.WriteTx, root []byte) error {
1169 if !t.editable() {
1170 return ErrSnapshotNotEditable
1171 }
1172
1173 if root == nil {
1174 return fmt.Errorf("can not SetRoot with nil root")
1175 }
1176
1177 // check that the root exists in the db
1178 if !bytes.Equal(root, t.emptyHash) {
1179 if _, err := wTx.Get(root); err == ErrKeyNotFound {
1180 return fmt.Errorf("can not SetRoot with root %x, as it"+
1181 " does not exist in the db", root)
1182 } else if err != nil {
1183 return err
1184 }
1185 }
1186
1187 return wTx.Set(dbKeyRoot, root)
1188}
1189
1190// Snapshot returns a read-only copy of the Tree from the given root
1191func (t *Tree) Snapshot(fromRoot []byte) (*Tree, error) {
1192 // allow to define which root to use
1193 if fromRoot == nil {
1194 var err error
1195 fromRoot, err = t.Root()
1196 if err != nil {
1197 return nil, err
1198 }
1199 }
1200 rTx := t.db.ReadTx()
1201 defer rTx.Discard()
1202 // check that the root exists in the db
1203 if !bytes.Equal(fromRoot, t.emptyHash) {
1204 if _, err := rTx.Get(fromRoot); err == ErrKeyNotFound {
1205 return nil,
1206 fmt.Errorf("can not do a Snapshot with root %x,"+
1207 " as it does not exist in the db", fromRoot)
1208 } else if err != nil {
1209 return nil, err
1210 }
1211 }
1212
1213 return &Tree{
1214 db: t.db,
1215 maxLevels: t.maxLevels,
1216 snapshotRoot: fromRoot,
1217 emptyHash: t.emptyHash,
1218 hashFunction: t.hashFunction,
1219 dbg: t.dbg,
1220 }, nil
1221}
1222
1223// Iterate iterates through the full Tree, executing the given function on each
1224// node of the Tree.
1225func (t *Tree) Iterate(fromRoot []byte, f func([]byte, []byte)) error {
1226 rTx := t.db.ReadTx()
1227 defer rTx.Discard()
1228
1229 return t.IterateWithTx(rTx, fromRoot, f)
1230}
1231
1232// IterateWithTx does the same than the Iterate method, but allowing to pass
1233// the db.ReadTx that is used.
1234func (t *Tree) IterateWithTx(rTx db.ReadTx, fromRoot []byte, f func([]byte, []byte)) error {
1235 // allow to define which root to use
1236 if fromRoot == nil {
1237 var err error
1238 fromRoot, err = t.RootWithTx(rTx)
1239 if err != nil {
1240 return err
1241 }
1242 }
1243 return t.iter(rTx, fromRoot, f)
1244}
1245
1246// IterateWithStop does the same than Iterate, but with int for the current
1247// level, and a boolean parameter used by the passed function, is to indicate to
1248// stop iterating on the branch when the method returns 'true'.
1249func (t *Tree) IterateWithStop(fromRoot []byte, f func(int, []byte, []byte) bool) error {
1250 rTx := t.db.ReadTx()
1251 defer rTx.Discard()
1252
1253 // allow to define which root to use
1254 if fromRoot == nil {
1255 var err error
1256 fromRoot, err = t.RootWithTx(rTx)
1257 if err != nil {
1258 return err
1259 }
1260 }
1261 return t.iterWithStop(rTx, fromRoot, 0, f)
1262}
1263
1264// IterateWithStopWithTx does the same than the IterateWithStop method, but
1265// allowing to pass the db.ReadTx that is used.
1266func (t *Tree) IterateWithStopWithTx(rTx db.ReadTx, fromRoot []byte,
1267 f func(int, []byte, []byte) bool) error {
1268 // allow to define which root to use
1269 if fromRoot == nil {
1270 var err error
1271 fromRoot, err = t.RootWithTx(rTx)
1272 if err != nil {
1273 return err
1274 }
1275 }
1276 return t.iterWithStop(rTx, fromRoot, 0, f)
1277}
1278
1279func (t *Tree) iterWithStop(rTx db.ReadTx, k []byte, currLevel int,
1280 f func(int, []byte, []byte) bool) error {
1281 var v []byte
1282 var err error
1283 if bytes.Equal(k, t.emptyHash) {
1284 v = t.emptyHash
1285 } else {
1286 v, err = rTx.Get(k)
1287 if err != nil {
1288 return err
1289 }
1290 }
1291 currLevel++
1292
1293 switch v[0] {
1294 case PrefixValueEmpty:
1295 f(currLevel, k, v)
1296 case PrefixValueLeaf:
1297 f(currLevel, k, v)
1298 case PrefixValueIntermediate:
1299 stop := f(currLevel, k, v)
1300 if stop {
1301 return nil
1302 }
1303 l, r := ReadIntermediateChilds(v)
1304 if err = t.iterWithStop(rTx, l, currLevel, f); err != nil {
1305 return err
1306 }
1307 if err = t.iterWithStop(rTx, r, currLevel, f); err != nil {
1308 return err
1309 }
1310 default:
1311 return ErrInvalidValuePrefix
1312 }
1313 return nil
1314}
1315
1316func (t *Tree) iter(rTx db.ReadTx, k []byte, f func([]byte, []byte)) error {
1317 f2 := func(currLvl int, k, v []byte) bool {
1318 f(k, v)
1319 return false
1320 }
1321 return t.iterWithStop(rTx, k, 0, f2)
1322}
1323
1324// Dump exports all the Tree leafs in a byte array
1325func (t *Tree) Dump(fromRoot []byte) ([]byte, error) {
1326 return t.dump(fromRoot, nil)
1327}
1328
1329// DumpWriter exports all the Tree leafs writing the bytes in the given Writer
1330func (t *Tree) DumpWriter(fromRoot []byte, w *bufio.Writer) error {
1331 _, err := t.dump(fromRoot, w)
1332 return err
1333}
1334
1335// dump exports all the Tree leafs. If the given w is nil, it will return a
1336// byte array with the dump, if w contains a *bufio.Writer, it will write the
1337// dump in w.
1338// The format of the dump is the following:
1339// Dump length: [ N * (2+len(k+v)) ]. Where N is the number of key-values, and for each k+v:
1340// [ 1 byte | 1 byte | S bytes | len(v) bytes ]
1341// [ len(k) | len(v) | key | value ]
1342// Where S is the size of the output of the hash function used for the Tree.
1343func (t *Tree) dump(fromRoot []byte, w *bufio.Writer) ([]byte, error) {
1344 // allow to define which root to use
1345 if fromRoot == nil {
1346 var err error
1347 fromRoot, err = t.Root()
1348 if err != nil {
1349 return nil, err
1350 }
1351 }
1352
1353 // WARNING current encoding only supports key & values of 255 bytes each
1354 // (due using only 1 byte for the length headers).
1355 var b []byte
1356 var callbackErr error
1357 err := t.IterateWithStop(fromRoot, func(_ int, k, v []byte) bool {
1358 if v[0] != PrefixValueLeaf {
1359 return false
1360 }
1361 leafK, leafV := ReadLeafValue(v)
1362 kv := make([]byte, 2+len(leafK)+len(leafV))
1363 if len(leafK) > maxUint8 {
1364 callbackErr = fmt.Errorf("len(leafK) > %v", maxUint8)
1365 return true
1366 }
1367 kv[0] = byte(len(leafK))
1368 if len(leafV) > maxUint8 {
1369 callbackErr = fmt.Errorf("len(leafV) > %v", maxUint8)
1370 return true
1371 }
1372 kv[1] = byte(len(leafV))
1373 copy(kv[2:2+len(leafK)], leafK)
1374 copy(kv[2+len(leafK):], leafV)
1375
1376 if w == nil {
1377 b = append(b, kv...)
1378 } else {
1379 n, err := w.Write(kv)
1380 if err != nil {
1381 callbackErr = fmt.Errorf("dump: w.Write, %s", err)
1382 return true
1383 }
1384 if n != len(kv) {
1385 callbackErr = fmt.Errorf("dump: w.Write n!=len(kv), %s", err)
1386 return true
1387 }
1388 err = w.Flush()
1389 if err != nil {
1390 callbackErr = fmt.Errorf("dump: w.Flush, %s", err)
1391 return true
1392 }
1393 }
1394 return false
1395 })
1396 if callbackErr != nil {
1397 return nil, callbackErr
1398 }
1399 return b, err
1400}
1401
1402// ImportDump imports the leafs (that have been exported with the Dump method)
1403// in the Tree, reading them from the given byte array.
1404func (t *Tree) ImportDump(b []byte) error {
1405 bytesReader := bytes.NewReader(b)
1406 r := bufio.NewReader(bytesReader)
1407 return t.ImportDumpReader(r)
1408}
1409
1410// ImportDumpReader imports the leafs (that have been exported with the Dump
1411// method) in the Tree, reading them from the given reader.
1412func (t *Tree) ImportDumpReader(r *bufio.Reader) error {
1413 if !t.editable() {
1414 return ErrSnapshotNotEditable
1415 }
1416 root, err := t.Root()
1417 if err != nil {
1418 return err
1419 }
1420 if !bytes.Equal(root, t.emptyHash) {
1421 return ErrTreeNotEmpty
1422 }
1423
1424 var keys, values [][]byte
1425 for {
1426 l := make([]byte, 2)
1427 _, err = io.ReadFull(r, l)
1428 if err == io.EOF {
1429 break
1430 } else if err != nil {
1431 return err
1432 }
1433 k := make([]byte, l[0])
1434 _, err = io.ReadFull(r, k)
1435 if err != nil {
1436 return err
1437 }
1438 v := make([]byte, l[1])
1439 _, err = io.ReadFull(r, v)
1440 if err != nil {
1441 return err
1442 }
1443 keys = append(keys, k)
1444 values = append(values, v)
1445 }
1446 if _, err = t.AddBatch(keys, values); err != nil {
1447 return err
1448 }
1449 return nil
1450}
1451
1452// Graphviz iterates across the full tree to generate a string Graphviz
1453// representation of the tree and writes it to w
1454func (t *Tree) Graphviz(w io.Writer, fromRoot []byte) error {
1455 return t.GraphvizFirstNLevels(w, fromRoot, t.maxLevels)
1456}
1457
1458// GraphvizFirstNLevels iterates across the first NLevels of the tree to
1459// generate a string Graphviz representation of the first NLevels of the tree
1460// and writes it to w
1461func (t *Tree) GraphvizFirstNLevels(w io.Writer, fromRoot []byte, untilLvl int) error {
1462 fmt.Fprintf(w, `digraph hierarchy {
1463node [fontname=Monospace,fontsize=10,shape=box]
1464`)
1465
1466 rTx := t.db.ReadTx()
1467 defer rTx.Discard()
1468
1469 if fromRoot == nil {
1470 var err error
1471 fromRoot, err = t.RootWithTx(rTx)
1472 if err != nil {
1473 return err
1474 }
1475 }
1476
1477 nEmpties := 0
1478 err := t.iterWithStop(rTx, fromRoot, 0, func(currLvl int, k, v []byte) bool {
1479 if currLvl == untilLvl {
1480 return true // to stop the iter from going down
1481 }
1482 switch v[0] {
1483 case PrefixValueEmpty:
1484 case PrefixValueLeaf:
1485 fmt.Fprintf(w, "\"%v\" [style=filled];\n", hex.EncodeToString(k[:nChars]))
1486 // key & value from the leaf
1487 kB, vB := ReadLeafValue(v)
1488 fmt.Fprintf(w, "\"%v\" -> {\"k:%v\\nv:%v\"}\n",
1489 hex.EncodeToString(k[:nChars]), hex.EncodeToString(kB[:nChars]),
1490 hex.EncodeToString(vB[:nChars]))
1491 fmt.Fprintf(w, "\"k:%v\\nv:%v\" [style=dashed]\n",
1492 hex.EncodeToString(kB[:nChars]), hex.EncodeToString(vB[:nChars]))
1493 case PrefixValueIntermediate:
1494 l, r := ReadIntermediateChilds(v)
1495 lStr := hex.EncodeToString(l[:nChars])
1496 rStr := hex.EncodeToString(r[:nChars])
1497 eStr := ""
1498 if bytes.Equal(l, t.emptyHash) {
1499 lStr = fmt.Sprintf("empty%v", nEmpties)
1500 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
1501 lStr)
1502 nEmpties++
1503 }
1504 if bytes.Equal(r, t.emptyHash) {
1505 rStr = fmt.Sprintf("empty%v", nEmpties)
1506 eStr += fmt.Sprintf("\"%v\" [style=dashed,label=0];\n",
1507 rStr)
1508 nEmpties++
1509 }
1510 fmt.Fprintf(w, "\"%v\" -> {\"%v\" \"%v\"}\n", hex.EncodeToString(k[:nChars]),
1511 lStr, rStr)
1512 fmt.Fprint(w, eStr)
1513 default:
1514 }
1515 return false
1516 })
1517 fmt.Fprintf(w, "}\n")
1518 return err
1519}
1520
1521// PrintGraphviz prints the output of Tree.Graphviz
1522func (t *Tree) PrintGraphviz(fromRoot []byte) error {
1523 if fromRoot == nil {
1524 var err error
1525 fromRoot, err = t.Root()
1526 if err != nil {
1527 return err
1528 }
1529 }
1530 return t.PrintGraphvizFirstNLevels(fromRoot, t.maxLevels)
1531}
1532
1533// PrintGraphvizFirstNLevels prints the output of Tree.GraphvizFirstNLevels
1534func (t *Tree) PrintGraphvizFirstNLevels(fromRoot []byte, untilLvl int) error {
1535 if fromRoot == nil {
1536 var err error
1537 fromRoot, err = t.Root()
1538 if err != nil {
1539 return err
1540 }
1541 }
1542 w := bytes.NewBufferString("")
1543 fmt.Fprintf(w,
1544 "--------\nGraphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+":\n")
1545 err := t.GraphvizFirstNLevels(w, fromRoot, untilLvl)
1546 if err != nil {
1547 fmt.Println(w)
1548 return err
1549 }
1550 fmt.Fprintf(w,
1551 "End of Graphviz of the Tree with Root "+hex.EncodeToString(fromRoot)+"\n--------\n")
1552
1553 fmt.Println(w)
1554 return nil
1555}
1556
1557// TODO circom proofs
1558// TODO data structure for proofs (including root, key, value, siblings,
1559// hashFunction) + method to verify that data structure