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 kvs, invalids := combineInKVSet(aKvs, kvs)
314
315 // proceed with CASE A
316 sortKvs(kvs)
317
318 // cutPowerOfTwo, the excedent add it as normal Tree.Add
319 kvsP2, kvsNonP2 := cutPowerOfTwo(kvs)
320 var invalids2 []int
321 if nCPU > 1 {
322 invalids2, err = t.buildTreeBottomUp(nCPU, kvsP2)
323 if err != nil {
324 return nil, nil, err
325 }
326 } else {
327 invalids2, err = t.buildTreeBottomUpSingleThread(kvsP2)
328 if err != nil {
329 return nil, nil, err
330 }
331 }
332 invalids = append(invalids, invalids2...)
333 // return the excedents which will be added at the full tree at the end
334 return invalids, kvsNonP2, nil
335}
336
337func (t *Tree) caseC(nCPU, l int, keysAtL [][]byte, kvs []kv) ([]int, error) {
338 // 1. go down until level L (L=log2(nBuckets)): keysAtL
339
340 var excedents []kv
341 buckets := splitInBuckets(kvs, nCPU)
342
343 // 2. use keys at level L as roots of the subtrees under each one
344 excedentsInBucket := make([][]kv, nCPU)
345 subRoots := make([][]byte, nCPU)
346 txs := make([]db.Tx, nCPU)
347 var wg sync.WaitGroup
348 wg.Add(nCPU)
349 for i := 0; i < nCPU; i++ {
350 go func(cpu int) {
351 var err error
352 txs[cpu], err = t.db.NewTx()
353 if err != nil {
354 panic(err) // TODO WIP
355 }
356 bucketTree := Tree{tx: txs[cpu], db: t.db, maxLevels: t.maxLevels,
357 hashFunction: t.hashFunction, root: keysAtL[cpu]}
358
359 // 3. do CASE B (with 1 cpu) for each key at level L
360 _, bucketExcedents, err := bucketTree.caseB(1, l, buckets[cpu])
361 if err != nil {
362 panic(err) // TODO WIP
363 // return nil, err
364 }
365 excedentsInBucket[cpu] = bucketExcedents
366 subRoots[cpu] = bucketTree.root
367 wg.Done()
368 }(i)
369 }
370 wg.Wait()
371
372 // merge buckets txs into Tree.tx
373 for i := 0; i < len(txs); i++ {
374 if err := t.tx.Add(txs[i]); err != nil {
375 return nil, err
376 }
377 }
378 for i := 0; i < len(excedentsInBucket); i++ {
379 excedents = append(excedents, excedentsInBucket[i]...)
380 }
381
382 // 4. go upFromKeys from the new roots of the subtrees
383 newRoot, err := t.upFromKeys(subRoots)
384 if err != nil {
385 return nil, err
386 }
387 t.root = newRoot
388
389 // add the key-values that have not been used yet
390 var invalids []int
391 for i := 0; i < len(excedents); i++ {
392 if err = t.add(0, excedents[i].k, excedents[i].v); err != nil {
393 invalids = append(invalids, excedents[i].pos)
394 }
395 }
396 return invalids, nil
397}
398
399func (t *Tree) caseD(nCPU, l int, keysAtL [][]byte, kvs []kv) ([]int, error) {
400 if nCPU == 1 { // CASE D, but with 1 cpu
401 var invalids []int
402 for i := 0; i < len(kvs); i++ {
403 if err := t.add(0, kvs[i].k, kvs[i].v); err != nil {
404 invalids = append(invalids, kvs[i].pos)
405 }
406 }
407 return invalids, nil
408 }
409
410 buckets := splitInBuckets(kvs, nCPU)
411
412 subRoots := make([][]byte, nCPU)
413 invalidsInBucket := make([][]int, nCPU)
414 txs := make([]db.Tx, nCPU)
415
416 var wg sync.WaitGroup
417 wg.Add(nCPU)
418 for i := 0; i < nCPU; i++ {
419 go func(cpu int) {
420 var err error
421 txs[cpu], err = t.db.NewTx()
422 if err != nil {
423 panic(err) // TODO WIP
424 }
425 // put already existing tx into txs[cpu], as txs[cpu]
426 // needs the pending key-values that are not in tree.db,
427 // but are in tree.tx
428 if err := txs[cpu].Add(t.tx); err != nil {
429 panic(err) // TODO WIP
430 }
431
432 bucketTree := Tree{tx: txs[cpu], db: t.db, maxLevels: t.maxLevels - l,
433 hashFunction: t.hashFunction, root: keysAtL[cpu]}
434
435 for j := 0; j < len(buckets[cpu]); j++ {
436 if err = bucketTree.add(l, buckets[cpu][j].k, buckets[cpu][j].v); err != nil {
437 invalidsInBucket[cpu] = append(invalidsInBucket[cpu], buckets[cpu][j].pos)
438 }
439 }
440 subRoots[cpu] = bucketTree.root
441 wg.Done()
442 }(i)
443 }
444 wg.Wait()
445
446 // merge buckets txs into Tree.tx
447 for i := 0; i < len(txs); i++ {
448 if err := t.tx.Add(txs[i]); err != nil {
449 return nil, err
450 }
451 }
452
453 newRoot, err := t.upFromKeys(subRoots)
454 if err != nil {
455 return nil, err
456 }
457 t.root = newRoot
458
459 var invalids []int
460 for i := 0; i < len(invalidsInBucket); i++ {
461 invalids = append(invalids, invalidsInBucket[i]...)
462 }
463
464 return invalids, nil
465}
466
467func splitInBuckets(kvs []kv, nBuckets int) [][]kv {
468 buckets := make([][]kv, nBuckets)
469 // 1. classify the keyvalues into buckets
470 for i := 0; i < len(kvs); i++ {
471 pair := kvs[i]
472
473 // bucketnum := keyToBucket(pair.k, nBuckets)
474 bucketnum := keyToBucket(pair.keyPath, nBuckets)
475 buckets[bucketnum] = append(buckets[bucketnum], pair)
476 }
477 return buckets
478}
479
480// TODO rename in a more 'real' name (calculate bucket from/for key)
481func keyToBucket(k []byte, nBuckets int) int {
482 nLevels := int(math.Log2(float64(nBuckets)))
483 b := make([]int, nBuckets)
484 for i := 0; i < nBuckets; i++ {
485 b[i] = i
486 }
487 r := b
488 mid := len(r) / 2 //nolint:gomnd
489 for i := 0; i < nLevels; i++ {
490 if int(k[i/8]&(1<<(i%8))) != 0 {
491 r = r[mid:]
492 mid = len(r) / 2 //nolint:gomnd
493 } else {
494 r = r[:mid]
495 mid = len(r) / 2 //nolint:gomnd
496 }
497 }
498 return r[0]
499}
500
501type kv struct {
502 pos int // original position in the array
503 keyPath []byte
504 k []byte
505 v []byte
506}
507
508// compareBytes compares byte slices where the bytes are compared from left to
509// right and each byte is compared by bit from right to left
510func compareBytes(a, b []byte) bool {
511 // WIP
512 for i := 0; i < len(a); i++ {
513 for j := 0; j < 8; j++ {
514 aBit := a[i] & (1 << j)
515 bBit := b[i] & (1 << j)
516 if aBit > bBit {
517 return false
518 } else if aBit < bBit {
519 return true
520 }
521 }
522 }
523 return false
524}
525
526// sortKvs sorts the kv by path
527func sortKvs(kvs []kv) {
528 sort.Slice(kvs, func(i, j int) bool {
529 return compareBytes(kvs[i].keyPath, kvs[j].keyPath)
530 })
531}
532
533func (t *Tree) keysValuesToKvs(ks, vs [][]byte) ([]kv, error) {
534 if len(ks) != len(vs) {
535 return nil, fmt.Errorf("len(keys)!=len(values) (%d!=%d)",
536 len(ks), len(vs))
537 }
538 kvs := make([]kv, len(ks))
539 for i := 0; i < len(ks); i++ {
540 keyPath := make([]byte, t.hashFunction.Len())
541 copy(keyPath[:], ks[i])
542 kvs[i].pos = i
543 kvs[i].keyPath = ks[i]
544 kvs[i].k = ks[i]
545 kvs[i].v = vs[i]
546 }
547
548 return kvs, nil
549}
550
551/*
552func (t *Tree) kvsToKeysValues(kvs []kv) ([][]byte, [][]byte) {
553 ks := make([][]byte, len(kvs))
554 vs := make([][]byte, len(kvs))
555 for i := 0; i < len(kvs); i++ {
556 ks[i] = kvs[i].k
557 vs[i] = kvs[i].v
558 }
559 return ks, vs
560}
561*/
562
563// buildTreeBottomUp splits the key-values into n Buckets (where n is the number
564// of CPUs), in parallel builds a subtree for each bucket, once all the subtrees
565// are built, uses the subtrees roots as keys for a new tree, which as result
566// will have the complete Tree build from bottom to up, where until the
567// log2(nCPU) level it has been computed in parallel.
568func (t *Tree) buildTreeBottomUp(nCPU int, kvs []kv) ([]int, error) {
569 buckets := splitInBuckets(kvs, nCPU)
570
571 subRoots := make([][]byte, nCPU)
572 invalidsInBucket := make([][]int, nCPU)
573 txs := make([]db.Tx, nCPU)
574
575 var wg sync.WaitGroup
576 wg.Add(nCPU)
577 for i := 0; i < nCPU; i++ {
578 go func(cpu int) {
579 sortKvs(buckets[cpu])
580
581 var err error
582 txs[cpu], err = t.db.NewTx()
583 if err != nil {
584 panic(err) // TODO
585 }
586 bucketTree := Tree{tx: txs[cpu], db: t.db, maxLevels: t.maxLevels,
587 hashFunction: t.hashFunction, root: t.emptyHash}
588
589 currInvalids, err := bucketTree.buildTreeBottomUpSingleThread(buckets[cpu])
590 if err != nil {
591 panic(err) // TODO
592 }
593 invalidsInBucket[cpu] = currInvalids
594 subRoots[cpu] = bucketTree.root
595 wg.Done()
596 }(i)
597 }
598 wg.Wait()
599
600 // merge buckets txs into Tree.tx
601 for i := 0; i < len(txs); i++ {
602 if err := t.tx.Add(txs[i]); err != nil {
603 return nil, err
604 }
605 }
606
607 newRoot, err := t.upFromKeys(subRoots)
608 if err != nil {
609 return nil, err
610 }
611 t.root = newRoot
612
613 var invalids []int
614 for i := 0; i < len(invalidsInBucket); i++ {
615 invalids = append(invalids, invalidsInBucket[i]...)
616 }
617 return invalids, err
618}
619
620// buildTreeBottomUpSingleThread builds the tree with the given []kv from bottom
621// to the root. keys & values must be sorted by path, and the array ks must be
622// length multiple of 2
623func (t *Tree) buildTreeBottomUpSingleThread(kvs []kv) ([]int, error) {
624 // TODO check that log2(len(leafs)) < t.maxLevels, if not, maxLevels
625 // would be reached and should return error
626
627 var invalids []int
628 // build the leafs
629 leafKeys := make([][]byte, len(kvs))
630 for i := 0; i < len(kvs); i++ {
631 // TODO handle the case where Key&Value == 0
632 leafKey, leafValue, err := newLeafValue(t.hashFunction, kvs[i].k, kvs[i].v)
633 if err != nil {
634 // return nil, err
635 invalids = append(invalids, kvs[i].pos)
636 }
637 // store leafKey & leafValue to db
638 if err := t.tx.Put(leafKey, leafValue); err != nil {
639 // return nil, err
640 invalids = append(invalids, kvs[i].pos)
641 }
642 leafKeys[i] = leafKey
643 }
644 r, err := t.upFromKeys(leafKeys)
645 if err != nil {
646 return invalids, err
647 }
648 t.root = r
649 return invalids, nil
650}
651
652// keys & values must be sorted by path, and the array ks must be length
653// multiple of 2
654func (t *Tree) upFromKeys(ks [][]byte) ([]byte, error) {
655 if len(ks) == 1 {
656 return ks[0], nil
657 }
658
659 var rKs [][]byte
660 for i := 0; i < len(ks); i += 2 {
661 // TODO handle the case where Key&Value == 0
662 k, v, err := newIntermediate(t.hashFunction, ks[i], ks[i+1])
663 if err != nil {
664 return nil, err
665 }
666 // store k-v to db
667 if err = t.tx.Put(k, v); err != nil {
668 return nil, err
669 }
670 rKs = append(rKs, k)
671 }
672 return t.upFromKeys(rKs)
673}
674
675func (t *Tree) getLeafs(root []byte) ([][]byte, [][]byte, error) {
676 var ks, vs [][]byte
677 err := t.iter(root, func(k, v []byte) {
678 if v[0] != PrefixValueLeaf {
679 return
680 }
681 leafK, leafV := readLeafValue(v)
682 ks = append(ks, leafK)
683 vs = append(vs, leafV)
684 })
685 return ks, vs, err
686}
687
688func (t *Tree) getKeysAtLevel(l int) ([][]byte, error) {
689 var keys [][]byte
690 err := t.iterWithStop(t.root, 0, func(currLvl int, k, v []byte) bool {
691 if currLvl == l && !bytes.Equal(k, t.emptyHash) {
692 keys = append(keys, k)
693 }
694 if currLvl >= l {
695 return true // to stop the iter from going down
696 }
697 return false
698 })
699
700 return keys, err
701}
702
703// cutPowerOfTwo returns []kv of length that is a power of 2, and a second []kv
704// with the extra elements that don't fit in a power of 2 length
705func cutPowerOfTwo(kvs []kv) ([]kv, []kv) {
706 x := len(kvs)
707 if (x & (x - 1)) != 0 {
708 p2 := highestPowerOfTwo(x)
709 return kvs[:p2], kvs[p2:]
710 }
711 return kvs, nil
712}
713
714func highestPowerOfTwo(n int) int {
715 res := 0
716 for i := n; i >= 1; i-- {
717 if (i & (i - 1)) == 0 {
718 res = i
719 break
720 }
721 }
722 return res
723}
724
725// combineInKVSet combines two kv array in one single array without repeated
726// keys.
727func combineInKVSet(base, toAdd []kv) ([]kv, []int) {
728 // TODO this is a naive version, this will be implemented in a more
729 // efficient way or through maps, or through sorted binary search
730 r := base
731 var invalids []int
732 for i := 0; i < len(toAdd); i++ {
733 e := false
734 // check if toAdd[i] exists in the base set
735 for j := 0; j < len(base); j++ {
736 if bytes.Equal(toAdd[i].k, base[j].k) {
737 e = true
738 }
739 }
740 if !e {
741 r = append(r, toAdd[i])
742 } else {
743 invalids = append(invalids, toAdd[i].pos)
744 }
745 }
746 return r, invalids
747}
748
749// func computeSimpleAddCost(nLeafs int) int {
750// // nLvls 2^nLvls
751// nLvls := int(math.Log2(float64(nLeafs)))
752// return nLvls * int(math.Pow(2, float64(nLvls)))
753// }
754//
755// func computeBottomUpAddCost(nLeafs int) int {
756// // 2^nLvls * 2 - 1
757// nLvls := int(math.Log2(float64(nLeafs)))
758// return (int(math.Pow(2, float64(nLvls))) * 2) - 1
759// }