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