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