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