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