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