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