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