arbo/addbatch_test.go
addbatch_test.goRaw
1package arbo
2
3import (
4 "encoding/hex"
5 "fmt"
6 "math/big"
7 "testing"
8 "time"
9
10 qt "github.com/frankban/quicktest"
11 "github.com/iden3/go-merkletree/db/memory"
12)
13
14func TestBatchAux(t *testing.T) { // TODO TMP this test will be delted
15 c := qt.New(t)
16
17 nLeafs := 16
18
19 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
20 c.Assert(err, qt.IsNil)
21 defer tree.db.Close()
22
23 start := time.Now()
24 for i := 0; i < nLeafs; i++ {
25 k := BigIntToBytes(big.NewInt(int64(i)))
26 v := BigIntToBytes(big.NewInt(int64(i * 2)))
27 if err := tree.Add(k, v); err != nil {
28 t.Fatal(err)
29 }
30 }
31 fmt.Println(time.Since(start))
32
33 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
34 c.Assert(err, qt.IsNil)
35 defer tree2.db.Close()
36
37 for i := 0; i < 8; i++ {
38 k := BigIntToBytes(big.NewInt(int64(i)))
39 v := BigIntToBytes(big.NewInt(int64(i * 2)))
40 if err := tree2.Add(k, v); err != nil {
41 t.Fatal(err)
42 }
43 }
44 // tree.PrintGraphviz(nil)
45 // tree2.PrintGraphviz(nil)
46
47 var keys, values [][]byte
48 for i := 8; i < nLeafs; i++ {
49 k := BigIntToBytes(big.NewInt(int64(i)))
50 v := BigIntToBytes(big.NewInt(int64(i * 2)))
51 keys = append(keys, k)
52 values = append(values, v)
53 }
54 start = time.Now()
55 indexes, err := tree2.AddBatchOpt(keys, values)
56 c.Assert(err, qt.IsNil)
57 fmt.Println(time.Since(start))
58 c.Check(len(indexes), qt.Equals, 0)
59
60 // check that both trees roots are equal
61 c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
62}
63
64func TestAddBatchCaseA(t *testing.T) {
65 c := qt.New(t)
66
67 nLeafs := 1024
68
69 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
70 c.Assert(err, qt.IsNil)
71 defer tree.db.Close()
72
73 start := time.Now()
74 for i := 0; i < nLeafs; i++ {
75 k := BigIntToBytes(big.NewInt(int64(i)))
76 v := BigIntToBytes(big.NewInt(int64(i * 2)))
77 if err := tree.Add(k, v); err != nil {
78 t.Fatal(err)
79 }
80 }
81 fmt.Println("time elapsed without CASE A: ", time.Since(start))
82
83 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
84 c.Assert(err, qt.IsNil)
85 defer tree2.db.Close()
86
87 var keys, values [][]byte
88 for i := 0; i < nLeafs; i++ {
89 k := BigIntToBytes(big.NewInt(int64(i)))
90 v := BigIntToBytes(big.NewInt(int64(i * 2)))
91 keys = append(keys, k)
92 values = append(values, v)
93 }
94 start = time.Now()
95 indexes, err := tree2.AddBatchOpt(keys, values)
96 c.Assert(err, qt.IsNil)
97 fmt.Println("time elapsed with CASE A: ", time.Since(start))
98 c.Check(len(indexes), qt.Equals, 0)
99
100 // check that both trees roots are equal
101 c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
102}
103
104func TestAddBatchCaseANotPowerOf2(t *testing.T) {
105 c := qt.New(t)
106
107 nLeafs := 1027
108
109 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
110 c.Assert(err, qt.IsNil)
111 defer tree.db.Close()
112
113 for i := 0; i < nLeafs; i++ {
114 k := BigIntToBytes(big.NewInt(int64(i)))
115 v := BigIntToBytes(big.NewInt(int64(i * 2)))
116 if err := tree.Add(k, v); err != nil {
117 t.Fatal(err)
118 }
119 }
120
121 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
122 c.Assert(err, qt.IsNil)
123 defer tree2.db.Close()
124
125 var keys, values [][]byte
126 for i := 0; i < nLeafs; i++ {
127 k := BigIntToBytes(big.NewInt(int64(i)))
128 v := BigIntToBytes(big.NewInt(int64(i * 2)))
129 keys = append(keys, k)
130 values = append(values, v)
131 }
132 indexes, err := tree2.AddBatchOpt(keys, values)
133 c.Assert(err, qt.IsNil)
134 c.Check(len(indexes), qt.Equals, 0)
135
136 // check that both trees roots are equal
137 c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
138}
139
140func TestAddBatchCaseB(t *testing.T) {
141 c := qt.New(t)
142
143 nLeafs := 1024
144
145 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
146 c.Assert(err, qt.IsNil)
147 defer tree.db.Close()
148
149 start := time.Now()
150 for i := 0; i < nLeafs; i++ {
151 k := BigIntToBytes(big.NewInt(int64(i)))
152 v := BigIntToBytes(big.NewInt(int64(i * 2)))
153 if err := tree.Add(k, v); err != nil {
154 t.Fatal(err)
155 }
156 }
157 fmt.Println("time elapsed without CASE B: ", time.Since(start))
158
159 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
160 c.Assert(err, qt.IsNil)
161 defer tree2.db.Close()
162
163 // add the initial leafs to fill a bit the tree before calling the
164 // AddBatch method
165 for i := 0; i < 99; i++ { // TMP TODO use const minLeafsThreshold-1 once ready
166 k := BigIntToBytes(big.NewInt(int64(i)))
167 v := BigIntToBytes(big.NewInt(int64(i * 2)))
168 if err := tree2.Add(k, v); err != nil {
169 t.Fatal(err)
170 }
171 }
172
173 var keys, values [][]byte
174 for i := 99; i < nLeafs; i++ {
175 k := BigIntToBytes(big.NewInt(int64(i)))
176 v := BigIntToBytes(big.NewInt(int64(i * 2)))
177 keys = append(keys, k)
178 values = append(values, v)
179 }
180 start = time.Now()
181 indexes, err := tree2.AddBatchOpt(keys, values)
182 c.Assert(err, qt.IsNil)
183 fmt.Println("time elapsed with CASE B: ", time.Since(start))
184 c.Check(len(indexes), qt.Equals, 0)
185
186 // check that both trees roots are equal
187 c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
188}
189
190func TestGetKeysAtLevel(t *testing.T) {
191 c := qt.New(t)
192
193 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
194 c.Assert(err, qt.IsNil)
195 defer tree.db.Close()
196
197 for i := 0; i < 32; i++ {
198 k := BigIntToBytes(big.NewInt(int64(i)))
199 v := BigIntToBytes(big.NewInt(int64(i * 2)))
200 if err := tree.Add(k, v); err != nil {
201 t.Fatal(err)
202 }
203 }
204
205 keys, err := tree.getKeysAtLevel(2)
206 c.Assert(err, qt.IsNil)
207 expected := []string{
208 "a5d5f14fce7348e40751496cf25d107d91b0bd043435b9577d778a01f8aa6111",
209 "e9e8dd9b28a7f81d1ff34cb5cefc0146dd848b31031a427b79bdadb62e7f6910",
210 }
211 for i := 0; i < len(keys); i++ {
212 c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
213 }
214
215 keys, err = tree.getKeysAtLevel(3)
216 c.Assert(err, qt.IsNil)
217 expected = []string{
218 "9f12c13e52bca96ad4882a26558e48ab67ddd63e062b839207e893d961390f01",
219 "16d246dd6826ec7346c7328f11c4261facf82d4689f33263ff6e207956a77f21",
220 "4a22cc901c6337daa17a431fa20170684b710e5f551509511492ec24e81a8f2f",
221 "470d61abcbd154977bffc9a9ec5a8daff0caabcf2a25e8441f604c79daa0f82d",
222 }
223 for i := 0; i < len(keys); i++ {
224 c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
225 }
226
227 keys, err = tree.getKeysAtLevel(4)
228 c.Assert(err, qt.IsNil)
229 expected = []string{
230 "7a5d1c81f7b96318012de3417e53d4f13df5b1337718651cd29d0cb0a66edd20",
231 "3408213e4e844bdf3355eb8781c74e31626812898c2dbe141ed6d2c92256fc1c",
232 "dfd8a4d0b6954a3e9f3892e655b58d456eeedf9367f27dfdd9bc2dd6a5577312",
233 "9e99fbec06fb2a6725997c12c4995f62725eb4cce4808523a5a5e80cca64b007",
234 "0befa1e070231dbf4e8ff841c05878cdec823e0c09594c24910a248b3ff5a628",
235 "b7131b0a15c772a57005a4dc5d0d6dd4b3414f5d9ee7408ce5e86c5ab3520e04",
236 "6d1abe0364077846a56bab1deb1a04883eb796b74fe531a7676a9a370f83ab21",
237 "4270116394bede69cf9cd72069eca018238080380bef5de75be8dcbbe968e105",
238 }
239 for i := 0; i < len(keys); i++ {
240 c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
241 }
242}
243
244func TestSplitInBuckets(t *testing.T) {
245 c := qt.New(t)
246
247 nLeafs := 16
248 kvs := make([]kv, nLeafs)
249 for i := 0; i < nLeafs; i++ {
250 k := BigIntToBytes(big.NewInt(int64(i)))
251 v := BigIntToBytes(big.NewInt(int64(i * 2)))
252 keyPath := make([]byte, 32)
253 copy(keyPath[:], k)
254 kvs[i].pos = i
255 kvs[i].keyPath = k
256 kvs[i].k = k
257 kvs[i].v = v
258 }
259
260 // check keyToBucket results for 4 buckets & 8 keys
261 c.Assert(keyToBucket(kvs[0].k, 4), qt.Equals, 0)
262 c.Assert(keyToBucket(kvs[1].k, 4), qt.Equals, 2)
263 c.Assert(keyToBucket(kvs[2].k, 4), qt.Equals, 1)
264 c.Assert(keyToBucket(kvs[3].k, 4), qt.Equals, 3)
265 c.Assert(keyToBucket(kvs[4].k, 4), qt.Equals, 0)
266 c.Assert(keyToBucket(kvs[5].k, 4), qt.Equals, 2)
267 c.Assert(keyToBucket(kvs[6].k, 4), qt.Equals, 1)
268 c.Assert(keyToBucket(kvs[7].k, 4), qt.Equals, 3)
269
270 // check keyToBucket results for 8 buckets & 8 keys
271 c.Assert(keyToBucket(kvs[0].k, 8), qt.Equals, 0)
272 c.Assert(keyToBucket(kvs[1].k, 8), qt.Equals, 4)
273 c.Assert(keyToBucket(kvs[2].k, 8), qt.Equals, 2)
274 c.Assert(keyToBucket(kvs[3].k, 8), qt.Equals, 6)
275 c.Assert(keyToBucket(kvs[4].k, 8), qt.Equals, 1)
276 c.Assert(keyToBucket(kvs[5].k, 8), qt.Equals, 5)
277 c.Assert(keyToBucket(kvs[6].k, 8), qt.Equals, 3)
278 c.Assert(keyToBucket(kvs[7].k, 8), qt.Equals, 7)
279
280 buckets := splitInBuckets(kvs, 4)
281
282 expected := [][]string{
283 {
284 "00000000", // bucket 0
285 "08000000",
286 "04000000",
287 "0c000000",
288 },
289 {
290 "02000000", // bucket 1
291 "0a000000",
292 "06000000",
293 "0e000000",
294 },
295 {
296 "01000000", // bucket 2
297 "09000000",
298 "05000000",
299 "0d000000",
300 },
301 {
302 "03000000", // bucket 3
303 "0b000000",
304 "07000000",
305 "0f000000",
306 },
307 }
308
309 for i := 0; i < len(buckets); i++ {
310 sortKvs(buckets[i])
311 c.Assert(len(buckets[i]), qt.Equals, len(expected[i]))
312 for j := 0; j < len(buckets[i]); j++ {
313 c.Check(hex.EncodeToString(buckets[i][j].k[:4]), qt.Equals, expected[i][j])
314 }
315 }
316}
317
318func TestAddBatchCaseC(t *testing.T) {
319 c := qt.New(t)
320
321 nLeafs := 1024
322
323 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
324 c.Assert(err, qt.IsNil)
325 defer tree.db.Close()
326
327 start := time.Now()
328 for i := 0; i < nLeafs; i++ {
329 k := BigIntToBytes(big.NewInt(int64(i)))
330 v := BigIntToBytes(big.NewInt(int64(i * 2)))
331 if err := tree.Add(k, v); err != nil {
332 t.Fatal(err)
333 }
334 }
335 fmt.Println("time elapsed without CASE C: ", time.Since(start))
336
337 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
338 c.Assert(err, qt.IsNil)
339 defer tree2.db.Close()
340
341 // add the initial leafs to fill a bit the tree before calling the
342 // AddBatch method
343 for i := 0; i < 101; i++ { // TMP TODO use const minLeafsThreshold-1 once ready
344 k := BigIntToBytes(big.NewInt(int64(i)))
345 v := BigIntToBytes(big.NewInt(int64(i * 2)))
346 if err := tree2.Add(k, v); err != nil {
347 t.Fatal(err)
348 }
349 }
350
351 var keys, values [][]byte
352 for i := 101; i < nLeafs; i++ {
353 k := BigIntToBytes(big.NewInt(int64(i)))
354 v := BigIntToBytes(big.NewInt(int64(i * 2)))
355 keys = append(keys, k)
356 values = append(values, v)
357 }
358 start = time.Now()
359 indexes, err := tree2.AddBatchOpt(keys, values)
360 c.Assert(err, qt.IsNil)
361 fmt.Println("time elapsed with CASE C: ", time.Since(start))
362 c.Check(len(indexes), qt.Equals, 0)
363
364 // check that both trees roots are equal
365 c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
366}
367
368func TestAddBatchCaseD(t *testing.T) {
369 c := qt.New(t)
370
371 nLeafs := 4096
372
373 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
374 c.Assert(err, qt.IsNil)
375 defer tree.db.Close()
376
377 start := time.Now()
378 for i := 0; i < nLeafs; i++ {
379 k := BigIntToBytes(big.NewInt(int64(i)))
380 v := BigIntToBytes(big.NewInt(int64(i * 2)))
381 if err := tree.Add(k, v); err != nil {
382 t.Fatal(err)
383 }
384 }
385 fmt.Println("time elapsed without CASE D: ", time.Since(start))
386
387 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
388 c.Assert(err, qt.IsNil)
389 defer tree2.db.Close()
390
391 // add the initial leafs to fill a bit the tree before calling the
392 // AddBatch method
393 for i := 0; i < 900; i++ { // TMP TODO use const minLeafsThreshold+1 once ready
394 k := BigIntToBytes(big.NewInt(int64(i)))
395 v := BigIntToBytes(big.NewInt(int64(i * 2)))
396 if err := tree2.Add(k, v); err != nil {
397 t.Fatal(err)
398 }
399 }
400
401 var keys, values [][]byte
402 for i := 900; i < nLeafs; i++ {
403 k := BigIntToBytes(big.NewInt(int64(i)))
404 v := BigIntToBytes(big.NewInt(int64(i * 2)))
405 keys = append(keys, k)
406 values = append(values, v)
407 }
408 start = time.Now()
409 indexes, err := tree2.AddBatchOpt(keys, values)
410 c.Assert(err, qt.IsNil)
411 fmt.Println("time elapsed with CASE D: ", time.Since(start))
412 c.Check(len(indexes), qt.Equals, 0)
413
414 // check that both trees roots are equal
415 c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
416}
417
418func TestAddBatchCaseE(t *testing.T) {
419 c := qt.New(t)
420
421 nLeafs := 4096
422
423 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
424 c.Assert(err, qt.IsNil)
425 defer tree.db.Close()
426
427 start := time.Now()
428 for i := 0; i < nLeafs; i++ {
429 k := BigIntToBytes(big.NewInt(int64(i)))
430 v := BigIntToBytes(big.NewInt(int64(i * 2)))
431 if err := tree.Add(k, v); err != nil {
432 t.Fatal(err)
433 }
434 }
435 fmt.Println("time elapsed without CASE E: ", time.Since(start))
436
437 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
438 c.Assert(err, qt.IsNil)
439 defer tree2.db.Close()
440
441 var keys, values [][]byte
442 // add the initial leafs to fill a bit the tree before calling the
443 // AddBatch method
444 for i := 0; i < 900; i++ { // TMP TODO use const minLeafsThreshold+1 once ready
445 k := BigIntToBytes(big.NewInt(int64(i)))
446 v := BigIntToBytes(big.NewInt(int64(i * 2)))
447 // use only the keys of one bucket, store the not used ones for
448 // later
449 if i%4 != 0 {
450 keys = append(keys, k)
451 values = append(values, v)
452 continue
453 }
454 if err := tree2.Add(k, v); err != nil {
455 t.Fatal(err)
456 }
457 }
458
459 for i := 900; i < nLeafs; i++ {
460 k := BigIntToBytes(big.NewInt(int64(i)))
461 v := BigIntToBytes(big.NewInt(int64(i * 2)))
462 keys = append(keys, k)
463 values = append(values, v)
464 }
465 start = time.Now()
466 indexes, err := tree2.AddBatchOpt(keys, values)
467 c.Assert(err, qt.IsNil)
468 fmt.Println("time elapsed with CASE E: ", time.Since(start))
469 c.Check(len(indexes), qt.Equals, 0)
470
471 // check that both trees roots are equal
472 c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
473}
474
475func TestHighestPowerOfTwo(t *testing.T) {
476 c := qt.New(t)
477 c.Assert(highestPowerOfTwo(31), qt.Equals, 16)
478 c.Assert(highestPowerOfTwo(32), qt.Equals, 32)
479 c.Assert(highestPowerOfTwo(33), qt.Equals, 32)
480 c.Assert(highestPowerOfTwo(63), qt.Equals, 32)
481 c.Assert(highestPowerOfTwo(64), qt.Equals, 64)
482}
483
484// func printLeafs(name string, t *Tree) {
485// w := bytes.NewBufferString("")
486//
487// err := t.Iterate(func(k, v []byte) {
488// if v[0] != PrefixValueLeaf {
489// return
490// }
491// leafK, _ := readLeafValue(v)
492// fmt.Fprintf(w, hex.EncodeToString(leafK[:4])+"\n")
493// })
494// if err != nil {
495// panic(err)
496// }
497// err = ioutil.WriteFile(name, w.Bytes(), 0644)
498// if err != nil {
499// panic(err)
500// }
501//
502// }
503
504// func TestComputeCosts(t *testing.T) {
505// fmt.Println(computeSimpleAddCost(10))
506// fmt.Println(computeBottomUpAddCost(10))
507//
508// fmt.Println(computeSimpleAddCost(1024))
509// fmt.Println(computeBottomUpAddCost(1024))
510// }
511
512// TODO test tree with nLeafs > minLeafsThreshold, but that at level L, there is
513// less keys than nBuckets (so CASE C could be applied if first few leafs are
514// added to balance the tree)
515
516// TODO for Cases tests, add initial keys, do snapshot, and then measure time of
517// adding the rest of keys with loop over normal Add, and with AddBatch
518
519// TODO test adding batch with repeated keys in the batch