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// AddBatchOpt is the WIP implementation of the AddBatch method in a more
149// optimized approach.
150func (t *Tree) AddBatchOpt(keys, values [][]byte) ([]int, error) {
151 t.updateAccessTime()
152 t.Lock()
153 defer t.Unlock()
154
155 // when len(keyvalues) is not a power of 2, cut at the biggest power of
156 // 2 under the len(keys), add those 2**n key-values using the AddBatch
157 // approach, and then add the remaining key-values using tree.Add.
158
159 kvs, err := t.keysValuesToKvs(keys, values)
160 if err != nil {
161 return nil, err
162 }
163
164 t.tx, err = t.db.NewTx()
165 if err != nil {
166 return nil, err
167 }
168
169 // if nCPU is not a power of two, cut at the highest power of two under
170 // nCPU
171 nCPU := highestPowerOfTwo(runtime.NumCPU())
172 l := int(math.Log2(float64(nCPU)))
173 var invalids []int
174
175 // CASE A: if nLeafs==0 (root==0)
176 if bytes.Equal(t.root, t.emptyHash) {
177 invalids, err = t.caseA(nCPU, kvs)
178 if err != nil {
179 return nil, err
180 }
181
182 if err = t.finalizeAddBatch(); err != nil {
183 return nil, err
184 }
185 return invalids, nil
186 }
187
188 // CASE B: if nLeafs<nBuckets
189 nLeafs, err := t.GetNLeafs()
190 if err != nil {
191 return nil, err
192 }
193 if nLeafs < minLeafsThreshold { // CASE B
194 var excedents []kv
195 invalids, excedents, err = t.caseB(nCPU, 0, kvs)
196 if err != nil {
197 return nil, err
198 }
199 // add the excedents
200 for i := 0; i < len(excedents); i++ {
201 err = t.add(0, excedents[i].k, excedents[i].v)
202 if err != nil {
203 invalids = append(invalids, excedents[i].pos)
204 }
205 }
206
207 if err = t.finalizeAddBatch(); err != nil {
208 return nil, err
209 }
210 return invalids, nil
211 }
212
213 keysAtL, err := t.getKeysAtLevel(l + 1)
214 if err != nil {
215 return nil, err
216 }
217
218 // CASE C: if nLeafs>=minLeafsThreshold && (nLeafs/nBuckets) < minLeafsThreshold
219 // available parallelization, will need to be a power of 2 (2**n)
220 if nLeafs >= minLeafsThreshold &&
221 (nLeafs/nCPU) < minLeafsThreshold &&
222 len(keysAtL) == nCPU {
223 invalids, err = t.caseC(nCPU, l, keysAtL, kvs)
224 if err != nil {
225 return nil, err
226 }
227
228 if err = t.finalizeAddBatch(); err != nil {
229 return nil, err
230 }
231 return invalids, nil
232 }
233
234 // CASE E
235 if len(keysAtL) != nCPU {
236 // CASE E: add one key at each bucket, and then do CASE D
237 buckets := splitInBuckets(kvs, nCPU)
238 kvs = []kv{}
239 for i := 0; i < len(buckets); i++ {
240 // add one leaf of the bucket, if there is an error when
241 // adding the k-v, try to add the next one of the bucket
242 // (until one is added)
243 var inserted int
244 for j := 0; j < len(buckets[i]); j++ {
245 if err := t.add(0, buckets[i][j].k, buckets[i][j].v); err == nil {
246 inserted = j
247 break
248 }
249 }
250
251 // put the buckets elements except the inserted one
252 kvs = append(kvs, buckets[i][:inserted]...)
253 kvs = append(kvs, buckets[i][inserted+1:]...)
254 }
255 keysAtL, err = t.getKeysAtLevel(l + 1)
256 if err != nil {
257 return nil, err
258 }
259 }
260
261 // CASE D
262 if len(keysAtL) == nCPU { // enter in CASE D if len(keysAtL)=nCPU, if not, CASE E
263 invalidsCaseD, err := t.caseD(nCPU, l, keysAtL, kvs)
264 if err != nil {
265 return nil, err
266 }
267 invalids = append(invalids, invalidsCaseD...)
268
269 if err = t.finalizeAddBatch(); err != nil {
270 return nil, err
271 }
272 return invalids, nil
273 }
274
275 // TODO update NLeafs from DB
276
277 return nil, fmt.Errorf("UNIMPLEMENTED")
278}
279
280func (t *Tree) finalizeAddBatch() error {
281 // store root to db
282 if err := t.tx.Put(dbKeyRoot, t.root); err != nil {
283 return err
284 }
285 // commit db tx
286 if err := t.tx.Commit(); err != nil {
287 return err
288 }
289 return nil
290}
291
292func (t *Tree) caseA(nCPU int, kvs []kv) ([]int, error) {
293 // if len(kvs) is not a power of 2, cut at the bigger power
294 // of two under len(kvs), build the tree with that, and add
295 // later the excedents
296 kvsP2, kvsNonP2 := cutPowerOfTwo(kvs)
297 invalids, err := t.buildTreeBottomUp(nCPU, kvsP2)
298 if err != nil {
299 return nil, err
300 }
301 for i := 0; i < len(kvsNonP2); i++ {
302 if err = t.add(0, kvsNonP2[i].k, kvsNonP2[i].v); err != nil {
303 invalids = append(invalids, kvsNonP2[i].pos)
304 }
305 }
306 return invalids, nil
307}
308
309func (t *Tree) caseB(nCPU, l int, kvs []kv) ([]int, []kv, error) {
310 // get already existing keys
311 aKs, aVs, err := t.getLeafs(t.root)
312 if err != nil {
313 return nil, nil, err
314 }
315 aKvs, err := t.keysValuesToKvs(aKs, aVs)
316 if err != nil {
317 return nil, nil, err
318 }
319 // add already existing key-values to the inputted key-values
320 kvs = append(kvs, aKvs...)
321
322 // proceed with CASE A
323 sortKvs(kvs)
324
325 // cutPowerOfTwo, the excedent add it as normal Tree.Add
326 kvsP2, kvsNonP2 := cutPowerOfTwo(kvs)
327 var invalids []int
328 if nCPU > 1 {
329 invalids, err = t.buildTreeBottomUp(nCPU, kvsP2)
330 if err != nil {
331 return nil, nil, err
332 }
333 } else {
334 invalids, err = t.buildTreeBottomUpSingleThread(kvsP2)
335 if err != nil {
336 return nil, nil, err
337 }
338 }
339 // return the excedents which will be added at the full tree at the end
340 return invalids, kvsNonP2, nil
341}
342
343func (t *Tree) caseC(nCPU, l int, keysAtL [][]byte, kvs []kv) ([]int, error) {
344 // 1. go down until level L (L=log2(nBuckets)): keysAtL
345
346 var excedents []kv
347 buckets := splitInBuckets(kvs, nCPU)
348
349 // 2. use keys at level L as roots of the subtrees under each one
350 excedentsInBucket := make([][]kv, nCPU)
351 subRoots := make([][]byte, nCPU)
352 txs := make([]db.Tx, nCPU)
353 var wg sync.WaitGroup
354 wg.Add(nCPU)
355 for i := 0; i < nCPU; i++ {
356 go func(cpu int) {
357 var err error
358 txs[cpu], err = t.db.NewTx()
359 if err != nil {
360 panic(err) // TODO WIP
361 }
362 bucketTree := Tree{tx: txs[cpu], db: t.db, maxLevels: t.maxLevels,
363 hashFunction: t.hashFunction, root: keysAtL[cpu]}
364
365 // 3. do CASE B (with 1 cpu) for each key at level L
366 _, bucketExcedents, err := bucketTree.caseB(1, l, buckets[cpu])
367 if err != nil {
368 panic(err)
369 // return nil, err
370 }
371 excedentsInBucket[cpu] = bucketExcedents
372 subRoots[cpu] = bucketTree.root
373 wg.Done()
374 }(i)
375 }
376 wg.Wait()
377
378 // merge buckets txs into Tree.tx
379 for i := 0; i < len(txs); i++ {
380 if err := t.tx.Add(txs[i]); err != nil {
381 return nil, err
382 }
383 }
384 for i := 0; i < len(excedentsInBucket); i++ {
385 excedents = append(excedents, excedentsInBucket[i]...)
386 }
387
388 // 4. go upFromKeys from the new roots of the subtrees
389 newRoot, err := t.upFromKeys(subRoots)
390 if err != nil {
391 return nil, err
392 }
393 t.root = newRoot
394
395 // add the key-values that have not been used yet
396 var invalids []int
397 for i := 0; i < len(excedents); i++ {
398 // Add until the level L
399 if err = t.add(0, excedents[i].k, excedents[i].v); err != nil {
400 invalids = append(invalids, excedents[i].pos) // TODO WIP
401 }
402 }
403 return invalids, nil
404}
405
406func (t *Tree) caseD(nCPU, l int, keysAtL [][]byte, kvs []kv) ([]int, error) {
407 if nCPU == 1 { // CASE D, but with 1 cpu
408 var invalids []int
409 for i := 0; i < len(kvs); i++ {
410 if err := t.add(0, kvs[i].k, kvs[i].v); err != nil {
411 invalids = append(invalids, kvs[i].pos)
412 }
413 }
414 return invalids, nil
415 }
416
417 buckets := splitInBuckets(kvs, nCPU)
418
419 subRoots := make([][]byte, nCPU)
420 invalidsInBucket := make([][]int, nCPU)
421 txs := make([]db.Tx, nCPU)
422
423 var wg sync.WaitGroup
424 wg.Add(nCPU)
425 for i := 0; i < nCPU; i++ {
426 go func(cpu int) {
427 var err error
428 txs[cpu], err = t.db.NewTx()
429 if err != nil {
430 panic(err) // TODO WIP
431 }
432 // put already existing tx into txs[cpu], as txs[cpu]
433 // needs the pending key-values that are not in tree.db,
434 // but are in tree.tx
435 if err := txs[cpu].Add(t.tx); err != nil {
436 panic(err) // TODO WIP
437 }
438
439 bucketTree := Tree{tx: txs[cpu], db: t.db, maxLevels: t.maxLevels - l,
440 hashFunction: t.hashFunction, root: keysAtL[cpu]}
441
442 for j := 0; j < len(buckets[cpu]); j++ {
443 if err = bucketTree.add(l, buckets[cpu][j].k, buckets[cpu][j].v); err != nil {
444 invalidsInBucket[cpu] = append(invalidsInBucket[cpu], buckets[cpu][j].pos)
445 }
446 }
447 subRoots[cpu] = bucketTree.root
448 wg.Done()
449 }(i)
450 }
451 wg.Wait()
452
453 // merge buckets txs into Tree.tx
454 for i := 0; i < len(txs); i++ {
455 if err := t.tx.Add(txs[i]); err != nil {
456 return nil, err
457 }
458 }
459
460 newRoot, err := t.upFromKeys(subRoots)
461 if err != nil {
462 return nil, err
463 }
464 t.root = newRoot
465
466 var invalids []int
467 for i := 0; i < len(invalidsInBucket); i++ {
468 invalids = append(invalids, invalidsInBucket[i]...)
469 }
470
471 return invalids, nil
472}
473
474func splitInBuckets(kvs []kv, nBuckets int) [][]kv {
475 buckets := make([][]kv, nBuckets)
476 // 1. classify the keyvalues into buckets
477 for i := 0; i < len(kvs); i++ {
478 pair := kvs[i]
479
480 // bucketnum := keyToBucket(pair.k, nBuckets)
481 bucketnum := keyToBucket(pair.keyPath, nBuckets)
482 buckets[bucketnum] = append(buckets[bucketnum], pair)
483 }
484 return buckets
485}
486
487// TODO rename in a more 'real' name (calculate bucket from/for key)
488func keyToBucket(k []byte, nBuckets int) int {
489 nLevels := int(math.Log2(float64(nBuckets)))
490 b := make([]int, nBuckets)
491 for i := 0; i < nBuckets; i++ {
492 b[i] = i
493 }
494 r := b
495 mid := len(r) / 2 //nolint:gomnd
496 for i := 0; i < nLevels; i++ {
497 if int(k[i/8]&(1<<(i%8))) != 0 {
498 r = r[mid:]
499 mid = len(r) / 2 //nolint:gomnd
500 } else {
501 r = r[:mid]
502 mid = len(r) / 2 //nolint:gomnd
503 }
504 }
505 return r[0]
506}
507
508type kv struct {
509 pos int // original position in the array
510 keyPath []byte
511 k []byte
512 v []byte
513}
514
515// compareBytes compares byte slices where the bytes are compared from left to
516// right and each byte is compared by bit from right to left
517func compareBytes(a, b []byte) bool {
518 // WIP
519 for i := 0; i < len(a); i++ {
520 for j := 0; j < 8; j++ {
521 aBit := a[i] & (1 << j)
522 bBit := b[i] & (1 << j)
523 if aBit > bBit {
524 return false
525 } else if aBit < bBit {
526 return true
527 }
528 }
529 }
530 return false
531}
532
533// sortKvs sorts the kv by path
534func sortKvs(kvs []kv) {
535 sort.Slice(kvs, func(i, j int) bool {
536 return compareBytes(kvs[i].keyPath, kvs[j].keyPath)
537 })
538}
539
540func (t *Tree) keysValuesToKvs(ks, vs [][]byte) ([]kv, error) {
541 if len(ks) != len(vs) {
542 return nil, fmt.Errorf("len(keys)!=len(values) (%d!=%d)",
543 len(ks), len(vs))
544 }
545 kvs := make([]kv, len(ks))
546 for i := 0; i < len(ks); i++ {
547 keyPath := make([]byte, t.hashFunction.Len())
548 copy(keyPath[:], ks[i])
549 kvs[i].pos = i
550 kvs[i].keyPath = ks[i]
551 kvs[i].k = ks[i]
552 kvs[i].v = vs[i]
553 }
554
555 return kvs, nil
556}
557
558/*
559func (t *Tree) kvsToKeysValues(kvs []kv) ([][]byte, [][]byte) {
560 ks := make([][]byte, len(kvs))
561 vs := make([][]byte, len(kvs))
562 for i := 0; i < len(kvs); i++ {
563 ks[i] = kvs[i].k
564 vs[i] = kvs[i].v
565 }
566 return ks, vs
567}
568*/
569
570// buildTreeBottomUp splits the key-values into n Buckets (where n is the number
571// of CPUs), in parallel builds a subtree for each bucket, once all the subtrees
572// are built, uses the subtrees roots as keys for a new tree, which as result
573// will have the complete Tree build from bottom to up, where until the
574// log2(nCPU) level it has been computed in parallel.
575func (t *Tree) buildTreeBottomUp(nCPU int, kvs []kv) ([]int, error) {
576 buckets := splitInBuckets(kvs, nCPU)
577
578 subRoots := make([][]byte, nCPU)
579 invalidsInBucket := make([][]int, nCPU)
580 txs := make([]db.Tx, nCPU)
581
582 var wg sync.WaitGroup
583 wg.Add(nCPU)
584 for i := 0; i < nCPU; i++ {
585 go func(cpu int) {
586 sortKvs(buckets[cpu])
587
588 var err error
589 txs[cpu], err = t.db.NewTx()
590 if err != nil {
591 panic(err) // TODO
592 }
593 bucketTree := Tree{tx: txs[cpu], db: t.db, maxLevels: t.maxLevels,
594 hashFunction: t.hashFunction, root: t.emptyHash}
595
596 currInvalids, err := bucketTree.buildTreeBottomUpSingleThread(buckets[cpu])
597 if err != nil {
598 panic(err) // TODO
599 }
600 invalidsInBucket[cpu] = currInvalids
601 subRoots[cpu] = bucketTree.root
602 wg.Done()
603 }(i)
604 }
605 wg.Wait()
606
607 // merge buckets txs into Tree.tx
608 for i := 0; i < len(txs); i++ {
609 if err := t.tx.Add(txs[i]); err != nil {
610 return nil, err
611 }
612 }
613
614 newRoot, err := t.upFromKeys(subRoots)
615 if err != nil {
616 return nil, err
617 }
618 t.root = newRoot
619
620 var invalids []int
621 for i := 0; i < len(invalidsInBucket); i++ {
622 invalids = append(invalids, invalidsInBucket[i]...)
623 }
624 return invalids, err
625}
626
627// buildTreeBottomUpSingleThread builds the tree with the given []kv from bottom
628// to the root. keys & values must be sorted by path, and the array ks must be
629// length multiple of 2
630func (t *Tree) buildTreeBottomUpSingleThread(kvs []kv) ([]int, error) {
631 // TODO check that log2(len(leafs)) < t.maxLevels, if not, maxLevels
632 // would be reached and should return error
633
634 var invalids []int
635 // build the leafs
636 leafKeys := make([][]byte, len(kvs))
637 for i := 0; i < len(kvs); i++ {
638 // TODO handle the case where Key&Value == 0
639 leafKey, leafValue, err := newLeafValue(t.hashFunction, kvs[i].k, kvs[i].v)
640 if err != nil {
641 // return nil, err
642 invalids = append(invalids, kvs[i].pos)
643 }
644 // store leafKey & leafValue to db
645 if err := t.tx.Put(leafKey, leafValue); err != nil {
646 // return nil, err
647 invalids = append(invalids, kvs[i].pos)
648 }
649 leafKeys[i] = leafKey
650 }
651 r, err := t.upFromKeys(leafKeys)
652 if err != nil {
653 return invalids, err
654 }
655 t.root = r
656 return invalids, nil
657}
658
659// keys & values must be sorted by path, and the array ks must be length
660// multiple of 2
661func (t *Tree) upFromKeys(ks [][]byte) ([]byte, error) {
662 if len(ks) == 1 {
663 return ks[0], nil
664 }
665
666 var rKs [][]byte
667 for i := 0; i < len(ks); i += 2 {
668 // TODO handle the case where Key&Value == 0
669 k, v, err := newIntermediate(t.hashFunction, ks[i], ks[i+1])
670 if err != nil {
671 return nil, err
672 }
673 // store k-v to db
674 if err = t.tx.Put(k, v); err != nil {
675 return nil, err
676 }
677 rKs = append(rKs, k)
678 }
679 return t.upFromKeys(rKs)
680}
681
682func (t *Tree) getLeafs(root []byte) ([][]byte, [][]byte, error) {
683 var ks, vs [][]byte
684 err := t.iter(root, func(k, v []byte) {
685 if v[0] != PrefixValueLeaf {
686 return
687 }
688 leafK, leafV := readLeafValue(v)
689 ks = append(ks, leafK)
690 vs = append(vs, leafV)
691 })
692 return ks, vs, err
693}
694
695func (t *Tree) getKeysAtLevel(l int) ([][]byte, error) {
696 var keys [][]byte
697 err := t.iterWithStop(t.root, 0, func(currLvl int, k, v []byte) bool {
698 if currLvl == l && !bytes.Equal(k, t.emptyHash) {
699 keys = append(keys, k)
700 }
701 if currLvl >= l {
702 return true // to stop the iter from going down
703 }
704 return false
705 })
706
707 return keys, err
708}
709
710// cutPowerOfTwo returns []kv of length that is a power of 2, and a second []kv
711// with the extra elements that don't fit in a power of 2 length
712func cutPowerOfTwo(kvs []kv) ([]kv, []kv) {
713 x := len(kvs)
714 if (x & (x - 1)) != 0 {
715 p2 := highestPowerOfTwo(x)
716 return kvs[:p2], kvs[p2:]
717 }
718 return kvs, nil
719}
720
721func highestPowerOfTwo(n int) int {
722 res := 0
723 for i := n; i >= 1; i-- {
724 if (i & (i - 1)) == 0 {
725 res = i
726 break
727 }
728 }
729 return res
730}
731
732// func computeSimpleAddCost(nLeafs int) int {
733// // nLvls 2^nLvls
734// nLvls := int(math.Log2(float64(nLeafs)))
735// return nLvls * int(math.Pow(2, float64(nLvls)))
736// }
737//
738// func computeBottomUpAddCost(nLeafs int) int {
739// // 2^nLvls * 2 - 1
740// nLvls := int(math.Log2(float64(nLeafs)))
741// return (int(math.Pow(2, float64(nLvls))) * 2) - 1
742// }