arbo/addbatch.go
addbatch.goRaw
1package arbo
2
3import (
4 "bytes"
5 "fmt"
6 "math"
7 "runtime"
8 "sort"
9 "sync"
10
11 "github.com/iden3/go-merkletree/db"
12)
13
14/*
15
16AddBatch design
17===============
18
19
20CASE A: Empty Tree --> if tree is empty (root==0)
21=================================================
22- Build the full tree from bottom to top (from all the leaf to the root)
23
24
25CASE B: ALMOST CASE A, Almost empty Tree --> if Tree has numLeafs < minLeafsThreshold
26==============================================================================
27- Get the Leafs (key & value) (iterate the tree from the current root getting
28the leafs)
29- Create a new empty Tree
30- Do CASE A for the new Tree, giving the already existing key&values (leafs)
31from the original Tree + the new key&values to be added from the AddBatch call
32
33
34 R R
35 / \ / \
36 A * / \
37 / \ / \
38 B C * *
39 / | / \
40 / | / \
41 / | / \
42 L: A B G D
43 / \
44 / \
45 / \
46 C *
47 / \
48 / \
49 / \
50 ... ... (nLeafs < minLeafsThreshold)
51
52
53CASE C: ALMOST CASE B --> if Tree has few Leafs (but numLeafs>=minLeafsThreshold)
54==============================================================================
55- Use A, B, G, F as Roots of subtrees
56- Do CASE B for each subtree
57- Then go from L to the Root
58
59 R
60 / \
61 / \
62 / \
63 * *
64 / | / \
65 / | / \
66 / | / \
67L: A B G D
68 / \
69 / \
70 / \
71 C *
72 / \
73 / \
74 / \
75 ... ... (nLeafs >= minLeafsThreshold)
76
77
78
79CASE D: Already populated Tree
80==============================
81- Use A, B, C, D as subtree
82- Sort the Keys in Buckets that share the initial part of the path
83- For each subtree add there the new leafs
84
85 R
86 / \
87 / \
88 / \
89 * *
90 / | / \
91 / | / \
92 / | / \
93L: A B C D
94 /\ /\ / \ / \
95 ... ... ... ... ... ...
96
97
98CASE E: Already populated Tree Unbalanced
99=========================================
100- Need to fill M1 and M2, and then will be able to use CASE D
101 - Search for M1 & M2 in the inputed Keys
102 - Add M1 & M2 to the Tree
103 - From here can use CASE D
104
105 R
106 / \
107 / \
108 / \
109 * *
110 | \
111 | \
112 | \
113L: M1 * M2 * (where M1 and M2 are empty)
114 / | /
115 / | /
116 / | /
117 A * *
118 / \ | \
119 / \ | \
120 / \ | \
121 B * * C
122 / \ |\
123 ... ... | \
124 | \
125 D E
126
127
128
129Algorithm decision
130==================
131- if nLeafs==0 (root==0): CASE A
132- if nLeafs<minLeafsThreshold: CASE B
133- if nLeafs>=minLeafsThreshold && (nLeafs/nBuckets) < minLeafsThreshold: CASE C
134- else: CASE D & CASE E
135
136
137- Multiple tree.Add calls: O(n log n)
138 - Used in: cases A, B, C
139- Tree from bottom to top: O(log n)
140 - Used in: cases D, E
141
142*/
143
144const (
145 minLeafsThreshold = 100 // nolint:gomnd // TMP WIP this will be autocalculated
146)
147
148// AddBatch adds a batch of key-values to the Tree. Returns an array containing
149// the indexes of the keys failed to add.
150func (t *Tree) AddBatch(keys, values [][]byte) ([]int, error) {
151 // TODO: support vaules=nil
152 t.updateAccessTime()
153 t.Lock()
154 defer t.Unlock()
155
156 kvs, err := t.keysValuesToKvs(keys, values)
157 if err != nil {
158 return nil, err
159 }
160
161 t.tx, err = t.db.NewTx()
162 if err != nil {
163 return nil, err
164 }
165
166 // if nCPU is not a power of two, cut at the highest power of two under
167 // nCPU
168 nCPU := flp2(runtime.NumCPU())
169 l := int(math.Log2(float64(nCPU)))
170 var invalids []int
171
172 // CASE A: if nLeafs==0 (root==0)
173 if bytes.Equal(t.root, t.emptyHash) {
174 invalids, err = t.caseA(nCPU, kvs)
175 if err != nil {
176 return nil, err
177 }
178
179 return t.finalizeAddBatch(len(keys), invalids)
180 }
181
182 // CASE B: if nLeafs<nBuckets
183 nLeafs, err := t.GetNLeafs()
184 if err != nil {
185 return nil, err
186 }
187 if nLeafs < minLeafsThreshold { // CASE B
188 invalids, err = t.caseB(nCPU, 0, kvs)
189 if err != nil {
190 return nil, err
191 }
192
193 return t.finalizeAddBatch(len(keys), invalids)
194 }
195
196 keysAtL, err := t.getKeysAtLevel(l + 1)
197 if err != nil {
198 return nil, err
199 }
200
201 // CASE C: if nLeafs>=minLeafsThreshold && (nLeafs/nBuckets) < minLeafsThreshold
202 // available parallelization, will need to be a power of 2 (2**n)
203 if nLeafs >= minLeafsThreshold &&
204 (nLeafs/nCPU) < minLeafsThreshold &&
205 len(keysAtL) == nCPU {
206 invalids, err = t.caseC(nCPU, l, keysAtL, kvs)
207 if err != nil {
208 return nil, err
209 }
210
211 return t.finalizeAddBatch(len(keys), invalids)
212 }
213
214 // CASE E
215 if len(keysAtL) != nCPU {
216 // CASE E: add one key at each bucket, and then do CASE D
217 buckets := splitInBuckets(kvs, nCPU)
218 kvs = []kv{}
219 for i := 0; i < len(buckets); i++ {
220 // add one leaf of the bucket, if there is an error when
221 // adding the k-v, try to add the next one of the bucket
222 // (until one is added)
223 var inserted int
224 for j := 0; j < len(buckets[i]); j++ {
225 if err := t.add(0, buckets[i][j].k, buckets[i][j].v); err == nil {
226 inserted = j
227 break
228 }
229 }
230
231 // put the buckets elements except the inserted one
232 kvs = append(kvs, buckets[i][:inserted]...)
233 kvs = append(kvs, buckets[i][inserted+1:]...)
234 }
235 keysAtL, err = t.getKeysAtLevel(l + 1)
236 if err != nil {
237 return nil, err
238 }
239 }
240
241 // CASE D
242 if len(keysAtL) == nCPU { // enter in CASE D if len(keysAtL)=nCPU, if not, CASE E
243 invalidsCaseD, err := t.caseD(nCPU, l, keysAtL, kvs)
244 if err != nil {
245 return nil, err
246 }
247 invalids = append(invalids, invalidsCaseD...)
248
249 return t.finalizeAddBatch(len(keys), invalids)
250 }
251
252 return nil, fmt.Errorf("UNIMPLEMENTED")
253}
254
255func (t *Tree) finalizeAddBatch(nKeys int, invalids []int) ([]int, error) {
256 // store root to db
257 if err := t.dbPut(dbKeyRoot, t.root); err != nil {
258 return nil, err
259 }
260
261 // update nLeafs
262 if err := t.incNLeafs(nKeys - len(invalids)); err != nil {
263 return nil, err
264 }
265
266 // commit db tx
267 if err := t.tx.Commit(); err != nil {
268 return nil, err
269 }
270 return invalids, nil
271}
272
273func (t *Tree) caseA(nCPU int, kvs []kv) ([]int, error) {
274 invalids, err := t.buildTreeFromLeafs(nCPU, kvs)
275 if err != nil {
276 return nil, err
277 }
278 return invalids, nil
279}
280
281func (t *Tree) caseB(nCPU, l int, kvs []kv) ([]int, error) {
282 // get already existing keys
283 aKs, aVs, err := t.getLeafs(t.root)
284 if err != nil {
285 return nil, err
286 }
287 aKvs, err := t.keysValuesToKvs(aKs, aVs)
288 if err != nil {
289 return nil, err
290 }
291 // add already existing key-values to the inputted key-values
292 // kvs = append(kvs, aKvs...)
293 kvs, invalids := combineInKVSet(aKvs, kvs)
294
295 // proceed with CASE A
296 sortKvs(kvs)
297
298 var invalids2 []int
299 if nCPU > 1 {
300 invalids2, err = t.buildTreeFromLeafs(nCPU, kvs)
301 if err != nil {
302 return nil, err
303 }
304 } else {
305 invalids2, err = t.buildTreeFromLeafsSingleThread(l, kvs)
306 if err != nil {
307 return nil, err
308 }
309 }
310 invalids = append(invalids, invalids2...)
311
312 return invalids, nil
313}
314
315func (t *Tree) caseC(nCPU, l int, keysAtL [][]byte, kvs []kv) ([]int, error) {
316 // 1. go down until level L (L=log2(nBuckets)): keysAtL
317
318 var excedents []kv
319 buckets := splitInBuckets(kvs, nCPU)
320
321 // 2. use keys at level L as roots of the subtrees under each one
322 subRoots := make([][]byte, nCPU)
323 dbgStatsPerBucket := make([]*dbgStats, nCPU)
324 txs := make([]db.Tx, nCPU)
325 var wg sync.WaitGroup
326 wg.Add(nCPU)
327 for i := 0; i < nCPU; i++ {
328 go func(cpu int) {
329 var err error
330 txs[cpu], err = t.db.NewTx()
331 if err != nil {
332 panic(err) // TODO WIP
333 }
334 if err := txs[cpu].Add(t.tx); err != nil {
335 panic(err) // TODO
336 }
337 bucketTree := Tree{tx: txs[cpu], db: t.db, maxLevels: t.maxLevels,
338 hashFunction: t.hashFunction, root: keysAtL[cpu],
339 emptyHash: t.emptyHash, dbg: newDbgStats()}
340
341 // 3. do CASE B (with 1 cpu) for each key at level L
342 _, err = bucketTree.caseB(1, l, buckets[cpu]) // TODO handle invalids
343 if err != nil {
344 panic(err) // TODO WIP
345 // return nil, err
346 }
347 subRoots[cpu] = bucketTree.root
348 dbgStatsPerBucket[cpu] = bucketTree.dbg
349 wg.Done()
350 }(i)
351 }
352 wg.Wait()
353
354 // merge buckets txs into Tree.tx
355 for i := 0; i < len(txs); i++ {
356 if err := t.tx.Add(txs[i]); err != nil {
357 return nil, err
358 }
359 }
360
361 // 4. go upFromKeys from the new roots of the subtrees
362 newRoot, err := t.upFromKeys(subRoots)
363 if err != nil {
364 return nil, err
365 }
366 t.root = newRoot
367
368 // add the key-values that have not been used yet
369 var invalids []int
370 for i := 0; i < len(excedents); i++ {
371 if err = t.add(0, excedents[i].k, excedents[i].v); err != nil {
372 invalids = append(invalids, excedents[i].pos)
373 }
374 }
375
376 for i := 0; i < len(dbgStatsPerBucket); i++ {
377 t.dbg.add(dbgStatsPerBucket[i])
378 }
379
380 return invalids, nil
381}
382
383func (t *Tree) caseD(nCPU, l int, keysAtL [][]byte, kvs []kv) ([]int, error) {
384 if nCPU == 1 { // CASE D, but with 1 cpu
385 var invalids []int
386 for i := 0; i < len(kvs); i++ {
387 if err := t.add(0, kvs[i].k, kvs[i].v); err != nil {
388 invalids = append(invalids, kvs[i].pos)
389 }
390 }
391 return invalids, nil
392 }
393
394 buckets := splitInBuckets(kvs, nCPU)
395
396 subRoots := make([][]byte, nCPU)
397 invalidsInBucket := make([][]int, nCPU)
398 dbgStatsPerBucket := make([]*dbgStats, nCPU)
399 txs := make([]db.Tx, nCPU)
400
401 var wg sync.WaitGroup
402 wg.Add(nCPU)
403 for i := 0; i < nCPU; i++ {
404 go func(cpu int) {
405 var err error
406 txs[cpu], err = t.db.NewTx()
407 if err != nil {
408 panic(err) // TODO WIP
409 }
410 // put already existing tx into txs[cpu], as txs[cpu]
411 // needs the pending key-values that are not in tree.db,
412 // but are in tree.tx
413 if err := txs[cpu].Add(t.tx); err != nil {
414 panic(err) // TODO WIP
415 }
416
417 bucketTree := Tree{tx: txs[cpu], db: t.db, maxLevels: t.maxLevels - l,
418 hashFunction: t.hashFunction, root: keysAtL[cpu],
419 emptyHash: t.emptyHash, dbg: newDbgStats()}
420
421 for j := 0; j < len(buckets[cpu]); j++ {
422 if err = bucketTree.add(l, buckets[cpu][j].k, buckets[cpu][j].v); err != nil {
423 invalidsInBucket[cpu] = append(invalidsInBucket[cpu], buckets[cpu][j].pos)
424 }
425 }
426 subRoots[cpu] = bucketTree.root
427 dbgStatsPerBucket[cpu] = bucketTree.dbg
428 wg.Done()
429 }(i)
430 }
431 wg.Wait()
432
433 // merge buckets txs into Tree.tx
434 for i := 0; i < len(txs); i++ {
435 if err := t.tx.Add(txs[i]); err != nil {
436 return nil, err
437 }
438 }
439
440 newRoot, err := t.upFromKeys(subRoots)
441 if err != nil {
442 return nil, err
443 }
444 t.root = newRoot
445
446 var invalids []int
447 for i := 0; i < len(invalidsInBucket); i++ {
448 invalids = append(invalids, invalidsInBucket[i]...)
449 }
450
451 for i := 0; i < len(dbgStatsPerBucket); i++ {
452 t.dbg.add(dbgStatsPerBucket[i])
453 }
454
455 return invalids, nil
456}
457
458func splitInBuckets(kvs []kv, nBuckets int) [][]kv {
459 buckets := make([][]kv, nBuckets)
460 // 1. classify the keyvalues into buckets
461 for i := 0; i < len(kvs); i++ {
462 pair := kvs[i]
463
464 // bucketnum := keyToBucket(pair.k, nBuckets)
465 bucketnum := keyToBucket(pair.keyPath, nBuckets)
466 buckets[bucketnum] = append(buckets[bucketnum], pair)
467 }
468 return buckets
469}
470
471// TODO rename in a more 'real' name (calculate bucket from/for key)
472func keyToBucket(k []byte, nBuckets int) int {
473 nLevels := int(math.Log2(float64(nBuckets)))
474 b := make([]int, nBuckets)
475 for i := 0; i < nBuckets; i++ {
476 b[i] = i
477 }
478 r := b
479 mid := len(r) / 2 //nolint:gomnd
480 for i := 0; i < nLevels; i++ {
481 if int(k[i/8]&(1<<(i%8))) != 0 {
482 r = r[mid:]
483 mid = len(r) / 2 //nolint:gomnd
484 } else {
485 r = r[:mid]
486 mid = len(r) / 2 //nolint:gomnd
487 }
488 }
489 return r[0]
490}
491
492type kv struct {
493 pos int // original position in the array
494 keyPath []byte
495 k []byte
496 v []byte
497}
498
499// compareBytes compares byte slices where the bytes are compared from left to
500// right and each byte is compared by bit from right to left
501func compareBytes(a, b []byte) bool {
502 // WIP
503 for i := 0; i < len(a); i++ {
504 for j := 0; j < 8; j++ {
505 aBit := a[i] & (1 << j)
506 bBit := b[i] & (1 << j)
507 if aBit > bBit {
508 return false
509 } else if aBit < bBit {
510 return true
511 }
512 }
513 }
514 return false
515}
516
517// sortKvs sorts the kv by path
518func sortKvs(kvs []kv) {
519 sort.Slice(kvs, func(i, j int) bool {
520 return compareBytes(kvs[i].keyPath, kvs[j].keyPath)
521 })
522}
523
524func (t *Tree) keysValuesToKvs(ks, vs [][]byte) ([]kv, error) {
525 if len(ks) != len(vs) {
526 return nil, fmt.Errorf("len(keys)!=len(values) (%d!=%d)",
527 len(ks), len(vs))
528 }
529 kvs := make([]kv, len(ks))
530 for i := 0; i < len(ks); i++ {
531 keyPath := make([]byte, t.hashFunction.Len())
532 copy(keyPath[:], ks[i])
533 kvs[i].pos = i
534 kvs[i].keyPath = keyPath
535 kvs[i].k = ks[i]
536 kvs[i].v = vs[i]
537 }
538
539 return kvs, nil
540}
541
542/*
543func (t *Tree) kvsToKeysValues(kvs []kv) ([][]byte, [][]byte) {
544 ks := make([][]byte, len(kvs))
545 vs := make([][]byte, len(kvs))
546 for i := 0; i < len(kvs); i++ {
547 ks[i] = kvs[i].k
548 vs[i] = kvs[i].v
549 }
550 return ks, vs
551}
552*/
553
554// buildTreeFromLeafs splits the key-values into n Buckets (where n is the number
555// of CPUs), in parallel builds a subtree for each bucket, once all the subtrees
556// are built, uses the subtrees roots as keys for a new tree, which as result
557// will have the complete Tree build from bottom to up, where until the
558// log2(nCPU) level it has been computed in parallel.
559func (t *Tree) buildTreeFromLeafs(nCPU int, kvs []kv) ([]int, error) {
560 l := int(math.Log2(float64(nCPU)))
561 buckets := splitInBuckets(kvs, nCPU)
562
563 subRoots := make([][]byte, nCPU)
564 invalidsInBucket := make([][]int, nCPU)
565 dbgStatsPerBucket := make([]*dbgStats, nCPU)
566 txs := make([]db.Tx, nCPU)
567
568 var wg sync.WaitGroup
569 wg.Add(nCPU)
570 for i := 0; i < nCPU; i++ {
571 go func(cpu int) {
572 sortKvs(buckets[cpu])
573
574 var err error
575 txs[cpu], err = t.db.NewTx()
576 if err != nil {
577 panic(err) // TODO
578 }
579 if err := txs[cpu].Add(t.tx); err != nil {
580 panic(err) // TODO
581 }
582 bucketTree := Tree{tx: txs[cpu], db: t.db, maxLevels: t.maxLevels,
583 hashFunction: t.hashFunction, root: t.emptyHash,
584 emptyHash: t.emptyHash, dbg: newDbgStats()}
585
586 currInvalids, err := bucketTree.buildTreeFromLeafsSingleThread(l, buckets[cpu])
587 if err != nil {
588 panic(err) // TODO
589 }
590 invalidsInBucket[cpu] = currInvalids
591 subRoots[cpu] = bucketTree.root
592 dbgStatsPerBucket[cpu] = bucketTree.dbg
593 wg.Done()
594 }(i)
595 }
596 wg.Wait()
597
598 // merge buckets txs into Tree.tx
599 for i := 0; i < len(txs); i++ {
600 if err := t.tx.Add(txs[i]); err != nil {
601 return nil, err
602 }
603 }
604
605 newRoot, err := t.upFromKeys(subRoots)
606 if err != nil {
607 return nil, err
608 }
609 t.root = newRoot
610
611 var invalids []int
612 for i := 0; i < len(invalidsInBucket); i++ {
613 invalids = append(invalids, invalidsInBucket[i]...)
614 }
615
616 for i := 0; i < len(dbgStatsPerBucket); i++ {
617 t.dbg.add(dbgStatsPerBucket[i])
618 }
619
620 return invalids, err
621}
622
623// buildTreeFromLeafsSingleThread builds the tree with the given []kv from bottom
624// to the root
625func (t *Tree) buildTreeFromLeafsSingleThread(l int, kvsRaw []kv) ([]int, error) {
626 // TODO check that log2(len(leafs)) < t.maxLevels, if not, maxLevels
627 // would be reached and should return error
628 if len(kvsRaw) == 0 {
629 return nil, nil
630 }
631
632 vt := newVT(t.maxLevels, t.hashFunction)
633 if t.dbg != nil {
634 vt.params.dbg = newDbgStats()
635 }
636
637 for i := 0; i < len(kvsRaw); i++ {
638 if err := vt.add(l, kvsRaw[i].k, kvsRaw[i].v); err != nil {
639 return nil, err
640 }
641 }
642 pairs, err := vt.computeHashes()
643 if err != nil {
644 return nil, err
645 }
646
647 // store pairs in db
648 for i := 0; i < len(pairs); i++ {
649 if err := t.dbPut(pairs[i][0], pairs[i][1]); err != nil {
650 return nil, err
651 }
652 }
653 t.dbg.add(vt.params.dbg)
654
655 // set tree.root from the virtual tree root
656 t.root = vt.root.h
657
658 return nil, nil // TODO invalids
659}
660
661// keys & values must be sorted by path, and the array ks must be length
662// multiple of 2
663func (t *Tree) upFromKeys(ks [][]byte) ([]byte, error) {
664 if len(ks) == 1 {
665 return ks[0], nil
666 }
667
668 var rKs [][]byte
669 for i := 0; i < len(ks); i += 2 {
670 if bytes.Equal(ks[i], t.emptyHash) && bytes.Equal(ks[i+1], t.emptyHash) {
671 // when both sub keys are empty, the key is also empty
672 rKs = append(rKs, t.emptyHash)
673 continue
674 }
675 k, v, err := newIntermediate(t.hashFunction, ks[i], ks[i+1])
676 if err != nil {
677 return nil, err
678 }
679 // store k-v to db
680 if err = t.dbPut(k, v); err != nil {
681 return nil, err
682 }
683 rKs = append(rKs, k)
684 }
685 return t.upFromKeys(rKs)
686}
687
688func (t *Tree) getLeafs(root []byte) ([][]byte, [][]byte, error) {
689 var ks, vs [][]byte
690 err := t.iter(root, func(k, v []byte) {
691 if v[0] != PrefixValueLeaf {
692 return
693 }
694 leafK, leafV := ReadLeafValue(v)
695 ks = append(ks, leafK)
696 vs = append(vs, leafV)
697 })
698 return ks, vs, err
699}
700
701func (t *Tree) getKeysAtLevel(l int) ([][]byte, error) {
702 var keys [][]byte
703 err := t.iterWithStop(t.root, 0, func(currLvl int, k, v []byte) bool {
704 if currLvl == l && !bytes.Equal(k, t.emptyHash) {
705 keys = append(keys, k)
706 }
707 if currLvl >= l {
708 return true // to stop the iter from going down
709 }
710 return false
711 })
712
713 return keys, err
714}
715
716// flp2 computes the floor power of 2, the highest power of 2 under the given
717// value.
718func flp2(n int) int {
719 res := 0
720 for i := n; i >= 1; i-- {
721 if (i & (i - 1)) == 0 {
722 res = i
723 break
724 }
725 }
726 return res
727}
728
729// combineInKVSet combines two kv array in one single array without repeated
730// keys.
731func combineInKVSet(base, toAdd []kv) ([]kv, []int) {
732 // TODO this is a naive version, this will be implemented in a more
733 // efficient way or through maps, or through sorted binary search
734 r := base
735 var invalids []int
736 for i := 0; i < len(toAdd); i++ {
737 e := false
738 // check if toAdd[i] exists in the base set
739 for j := 0; j < len(base); j++ {
740 if bytes.Equal(toAdd[i].k, base[j].k) {
741 e = true
742 }
743 }
744 if !e {
745 r = append(r, toAdd[i])
746 } else {
747 invalids = append(invalids, toAdd[i].pos)
748 }
749 }
750 return r, invalids
751}
752
753// loadVT loads a new virtual tree (vt) from the current Tree, which contains
754// the same leafs.
755// func (t *Tree) loadVT() (vt, error) {
756// vt := newVT(t.maxLevels, t.hashFunction)
757// vt.params.dbg = t.dbg
758// err := t.Iterate(func(k, v []byte) {
759// switch v[0] {
760// case PrefixValueEmpty:
761// case PrefixValueLeaf:
762// leafK, leafV := ReadLeafValue(v)
763// if err := vt.add(0, leafK, leafV); err != nil {
764// panic(err)
765// }
766// case PrefixValueIntermediate:
767// default:
768// }
769// })
770//
771// return vt, err
772// }
773
774// func computeSimpleAddCost(nLeafs int) int {
775// // nLvls 2^nLvls
776// nLvls := int(math.Log2(float64(nLeafs)))
777// return nLvls * int(math.Pow(2, float64(nLvls)))
778// }
779//
780// func computeFromLeafsAddCost(nLeafs int) int {
781// // 2^nLvls * 2 - 1
782// nLvls := int(math.Log2(float64(nLeafs)))
783// return (int(math.Pow(2, float64(nLvls))) * 2) - 1
784// }