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