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