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