arbo/addbatch_test.go
addbatch_test.goRaw
1package arbo
2
3import (
4 "crypto/rand"
5 "encoding/hex"
6 "fmt"
7 "math/big"
8 "runtime"
9 "testing"
10 "time"
11
12 qt "github.com/frankban/quicktest"
13 "github.com/iden3/go-merkletree/db/leveldb"
14 "github.com/iden3/go-merkletree/db/memory"
15)
16
17var debug = false
18
19func printTestContext(prefix string, nLeafs int, hashName, dbName string) {
20 if debug {
21 fmt.Printf("%snCPU: %d, nLeafs: %d, hash: %s, db: %s\n",
22 prefix, runtime.NumCPU(), nLeafs, hashName, dbName)
23 }
24}
25
26func printRes(name string, duration time.Duration) {
27 if debug {
28 fmt.Printf("%s: %s \n", name, duration)
29 }
30}
31
32func debugTime(descr string, time1, time2 time.Duration) {
33 if debug {
34 fmt.Printf("%s was %f times faster than without AddBatch\n",
35 descr, float64(time1)/float64(time2))
36 }
37}
38
39func testInit(c *qt.C, n int) (*Tree, *Tree) {
40 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
41 c.Assert(err, qt.IsNil)
42 defer tree1.db.Close()
43
44 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
45 c.Assert(err, qt.IsNil)
46 defer tree2.db.Close()
47
48 // add the initial leafs to fill a bit the trees before calling the
49 // AddBatch method
50 for i := 0; i < n; i++ {
51 k := BigIntToBytes(big.NewInt(int64(i)))
52 v := BigIntToBytes(big.NewInt(int64(i * 2)))
53 if err := tree1.Add(k, v); err != nil {
54 c.Fatal(err)
55 }
56 if err := tree2.Add(k, v); err != nil {
57 c.Fatal(err)
58 }
59 }
60 return tree1, tree2
61}
62
63func TestAddBatchCaseA(t *testing.T) {
64 c := qt.New(t)
65
66 nLeafs := 1024
67
68 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
69 c.Assert(err, qt.IsNil)
70 defer tree.db.Close()
71
72 start := time.Now()
73 for i := 0; i < nLeafs; i++ {
74 k := BigIntToBytes(big.NewInt(int64(i)))
75 v := BigIntToBytes(big.NewInt(int64(i * 2)))
76 if err := tree.Add(k, v); err != nil {
77 t.Fatal(err)
78 }
79 }
80 time1 := time.Since(start)
81
82 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
83 c.Assert(err, qt.IsNil)
84 defer tree2.db.Close()
85 tree2.dbgInit()
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.AddBatch(keys, values)
96 c.Assert(err, qt.IsNil)
97 time2 := time.Since(start)
98 if debug {
99 debugTime("CASE A, AddBatch", time1, time2)
100 printTestContext(" ", nLeafs, "Poseidon", "memory")
101 tree2.dbg.print(" ")
102 }
103 c.Check(len(indexes), qt.Equals, 0)
104
105 // check that both trees roots are equal
106 c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
107}
108
109func TestAddBatchCaseANotPowerOf2(t *testing.T) {
110 c := qt.New(t)
111
112 nLeafs := 1027
113
114 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
115 c.Assert(err, qt.IsNil)
116 defer tree.db.Close()
117
118 for i := 0; i < nLeafs; i++ {
119 k := BigIntToBytes(big.NewInt(int64(i)))
120 v := BigIntToBytes(big.NewInt(int64(i * 2)))
121 if err := tree.Add(k, v); err != nil {
122 t.Fatal(err)
123 }
124 }
125
126 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
127 c.Assert(err, qt.IsNil)
128 defer tree2.db.Close()
129
130 var keys, values [][]byte
131 for i := 0; i < nLeafs; i++ {
132 k := BigIntToBytes(big.NewInt(int64(i)))
133 v := BigIntToBytes(big.NewInt(int64(i * 2)))
134 keys = append(keys, k)
135 values = append(values, v)
136 }
137 indexes, err := tree2.AddBatch(keys, values)
138 c.Assert(err, qt.IsNil)
139 c.Check(len(indexes), qt.Equals, 0)
140
141 // check that both trees roots are equal
142 c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
143}
144
145func randomBytes(n int) []byte {
146 b := make([]byte, n)
147 _, err := rand.Read(b)
148 if err != nil {
149 panic(err)
150 }
151 return b
152}
153
154func TestBuildTreeFromLeafsSingleThread(t *testing.T) {
155 c := qt.New(t)
156 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
157 c.Assert(err, qt.IsNil)
158 defer tree1.db.Close()
159
160 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
161 c.Assert(err, qt.IsNil)
162 defer tree2.db.Close()
163
164 testvectorKeys := []string{
165 "1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642",
166 "2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf",
167 "1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5",
168 "d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7",
169 }
170 var keys, values [][]byte
171 for i := 0; i < len(testvectorKeys); i++ {
172 key, err := hex.DecodeString(testvectorKeys[i])
173 c.Assert(err, qt.IsNil)
174 keys = append(keys, key)
175 values = append(values, []byte{0})
176 }
177
178 for i := 0; i < len(keys); i++ {
179 if err := tree1.Add(keys[i], values[i]); err != nil {
180 t.Fatal(err)
181 }
182 }
183
184 kvs, err := tree2.keysValuesToKvs(keys, values)
185 c.Assert(err, qt.IsNil)
186 sortKvs(kvs)
187
188 tree2.tx, err = tree2.db.NewTx()
189 c.Assert(err, qt.IsNil)
190 // indexes, err := tree2.buildTreeFromLeafsSingleThread(kvs)
191 indexes, err := tree2.buildTreeFromLeafs(4, kvs)
192 c.Assert(err, qt.IsNil)
193 // tree1.PrintGraphviz(nil)
194 // tree2.PrintGraphviz(nil)
195
196 c.Check(len(indexes), qt.Equals, 0)
197
198 // check that both trees roots are equal
199 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
200 // 15b6a23945ae6c81342b7eb14e70fff50812dc8791cb15ec791eb08f91784139
201}
202
203func TestAddBatchCaseATestVector(t *testing.T) {
204 c := qt.New(t)
205 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
206 c.Assert(err, qt.IsNil)
207 defer tree1.db.Close()
208
209 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
210 c.Assert(err, qt.IsNil)
211 defer tree2.db.Close()
212
213 // leafs in 2nd level subtrees: [ 6, 0, 1, 1]
214 testvectorKeys := []string{
215 "1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642",
216 "2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf",
217 "1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5",
218 "d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7",
219 }
220 var keys, values [][]byte
221 for i := 0; i < len(testvectorKeys); i++ {
222 key, err := hex.DecodeString(testvectorKeys[i])
223 c.Assert(err, qt.IsNil)
224 keys = append(keys, key)
225 values = append(values, []byte{0})
226 }
227
228 for i := 0; i < len(keys); i++ {
229 if err := tree1.Add(keys[i], values[i]); err != nil {
230 t.Fatal(err)
231 }
232 }
233
234 indexes, err := tree2.AddBatch(keys, values)
235 c.Assert(err, qt.IsNil)
236 // tree1.PrintGraphviz(nil)
237 // tree2.PrintGraphviz(nil)
238
239 c.Check(len(indexes), qt.Equals, 0)
240
241 // check that both trees roots are equal
242 // fmt.Println(hex.EncodeToString(tree1.Root()))
243 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
244}
245
246func TestAddBatchCaseARandomKeys(t *testing.T) {
247 c := qt.New(t)
248
249 nLeafs := 8
250
251 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
252 c.Assert(err, qt.IsNil)
253 defer tree1.db.Close()
254
255 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
256 c.Assert(err, qt.IsNil)
257 defer tree2.db.Close()
258
259 var keys, values [][]byte
260 for i := 0; i < nLeafs; i++ {
261 keys = append(keys, randomBytes(32))
262 // values = append(values, randomBytes(32))
263 values = append(values, []byte{0})
264 // fmt.Println("K", hex.EncodeToString(keys[i]))
265 }
266
267 // TMP:
268 keys[0], _ = hex.DecodeString("1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642")
269 keys[1], _ = hex.DecodeString("2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf")
270 keys[2], _ = hex.DecodeString("9cb87ec67e875c61390edcd1ab517f443591047709a4d4e45b0f9ed980857b8e")
271 keys[3], _ = hex.DecodeString("9b4e9e92e974a589f426ceeb4cb291dc24893513fecf8e8460992dcf52621d4d")
272 keys[4], _ = hex.DecodeString("1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5")
273 keys[5], _ = hex.DecodeString("d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7")
274 keys[6], _ = hex.DecodeString("3cd55dbfb8f975f20a0925dfbdabe79fa2d51dd0268afbb8ba6b01de9dfcdd3c")
275 keys[7], _ = hex.DecodeString("5d0a9d6d9f197c091bf054fac9cb60e11ec723d6610ed8578e617b4d46cb43d5")
276
277 for i := 0; i < len(keys); i++ {
278 if err := tree1.Add(keys[i], values[i]); err != nil {
279 t.Fatal(err)
280 }
281 }
282
283 indexes, err := tree2.AddBatch(keys, values)
284 c.Assert(err, qt.IsNil)
285 // tree1.PrintGraphviz(nil)
286 // tree2.PrintGraphviz(nil)
287
288 c.Check(len(indexes), qt.Equals, 0)
289
290 // check that both trees roots are equal
291 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
292}
293
294func TestAddBatchCaseB(t *testing.T) {
295 c := qt.New(t)
296
297 nLeafs := 1024
298 initialNLeafs := 99 // TMP TODO use const minLeafsThreshold-1 once ready
299
300 tree1, tree2 := testInit(c, initialNLeafs)
301 tree2.dbgInit()
302
303 start := time.Now()
304 for i := initialNLeafs; i < nLeafs; i++ {
305 k := BigIntToBytes(big.NewInt(int64(i)))
306 v := BigIntToBytes(big.NewInt(int64(i * 2)))
307 if err := tree1.Add(k, v); err != nil {
308 t.Fatal(err)
309 }
310 }
311 time1 := time.Since(start)
312
313 // prepare the key-values to be added
314 var keys, values [][]byte
315 for i := initialNLeafs; i < nLeafs; i++ {
316 k := BigIntToBytes(big.NewInt(int64(i)))
317 v := BigIntToBytes(big.NewInt(int64(i * 2)))
318 keys = append(keys, k)
319 values = append(values, v)
320 }
321 start = time.Now()
322 indexes, err := tree2.AddBatch(keys, values)
323 c.Assert(err, qt.IsNil)
324 time2 := time.Since(start)
325 if debug {
326 debugTime("CASE B, AddBatch", time1, time2)
327 printTestContext(" ", nLeafs, "Poseidon", "memory")
328 tree2.dbg.print(" ")
329 }
330 c.Check(len(indexes), qt.Equals, 0)
331
332 // check that both trees roots are equal
333 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
334}
335
336func TestAddBatchCaseBRepeatedLeafs(t *testing.T) {
337 c := qt.New(t)
338
339 nLeafs := 1024
340 initialNLeafs := 99 // TMP TODO use const minLeafsThreshold-1 once ready
341
342 tree1, tree2 := testInit(c, initialNLeafs)
343
344 for i := initialNLeafs; i < nLeafs; i++ {
345 k := BigIntToBytes(big.NewInt(int64(i)))
346 v := BigIntToBytes(big.NewInt(int64(i * 2)))
347 if err := tree1.Add(k, v); err != nil {
348 t.Fatal(err)
349 }
350 }
351
352 // prepare the key-values to be added, including already existing keys
353 var keys, values [][]byte
354 for i := 0; i < nLeafs; i++ {
355 k := BigIntToBytes(big.NewInt(int64(i)))
356 v := BigIntToBytes(big.NewInt(int64(i * 2)))
357 keys = append(keys, k)
358 values = append(values, v)
359 }
360 indexes, err := tree2.AddBatch(keys, values)
361 c.Assert(err, qt.IsNil)
362 c.Check(len(indexes), qt.Equals, initialNLeafs)
363
364 // check that both trees roots are equal
365 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
366}
367
368func TestCombineInKVSet(t *testing.T) {
369 c := qt.New(t)
370
371 var a, b, expected []kv
372 for i := 0; i < 10; i++ {
373 k := BigIntToBytes(big.NewInt(int64(i)))
374 kv := kv{k: k}
375 if i < 7 {
376 a = append(a, kv)
377 }
378 if i >= 4 {
379 b = append(b, kv)
380 }
381 expected = append(expected, kv)
382 }
383
384 r, invalids := combineInKVSet(a, b)
385 for i := 0; i < len(r); i++ {
386 c.Assert(r[i].k, qt.DeepEquals, expected[i].k)
387 }
388 c.Assert(len(invalids), qt.Equals, 7-4)
389}
390
391func TestGetKeysAtLevel(t *testing.T) {
392 c := qt.New(t)
393
394 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
395 c.Assert(err, qt.IsNil)
396 defer tree.db.Close()
397
398 for i := 0; i < 32; i++ {
399 k := BigIntToBytes(big.NewInt(int64(i)))
400 v := BigIntToBytes(big.NewInt(int64(i * 2)))
401 if err := tree.Add(k, v); err != nil {
402 t.Fatal(err)
403 }
404 }
405
406 keys, err := tree.getKeysAtLevel(2)
407 c.Assert(err, qt.IsNil)
408 expected := []string{
409 "a5d5f14fce7348e40751496cf25d107d91b0bd043435b9577d778a01f8aa6111",
410 "e9e8dd9b28a7f81d1ff34cb5cefc0146dd848b31031a427b79bdadb62e7f6910",
411 }
412 for i := 0; i < len(keys); i++ {
413 c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
414 }
415
416 keys, err = tree.getKeysAtLevel(3)
417 c.Assert(err, qt.IsNil)
418 expected = []string{
419 "9f12c13e52bca96ad4882a26558e48ab67ddd63e062b839207e893d961390f01",
420 "16d246dd6826ec7346c7328f11c4261facf82d4689f33263ff6e207956a77f21",
421 "4a22cc901c6337daa17a431fa20170684b710e5f551509511492ec24e81a8f2f",
422 "470d61abcbd154977bffc9a9ec5a8daff0caabcf2a25e8441f604c79daa0f82d",
423 }
424 for i := 0; i < len(keys); i++ {
425 c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
426 }
427
428 keys, err = tree.getKeysAtLevel(4)
429 c.Assert(err, qt.IsNil)
430 expected = []string{
431 "7a5d1c81f7b96318012de3417e53d4f13df5b1337718651cd29d0cb0a66edd20",
432 "3408213e4e844bdf3355eb8781c74e31626812898c2dbe141ed6d2c92256fc1c",
433 "dfd8a4d0b6954a3e9f3892e655b58d456eeedf9367f27dfdd9bc2dd6a5577312",
434 "9e99fbec06fb2a6725997c12c4995f62725eb4cce4808523a5a5e80cca64b007",
435 "0befa1e070231dbf4e8ff841c05878cdec823e0c09594c24910a248b3ff5a628",
436 "b7131b0a15c772a57005a4dc5d0d6dd4b3414f5d9ee7408ce5e86c5ab3520e04",
437 "6d1abe0364077846a56bab1deb1a04883eb796b74fe531a7676a9a370f83ab21",
438 "4270116394bede69cf9cd72069eca018238080380bef5de75be8dcbbe968e105",
439 }
440 for i := 0; i < len(keys); i++ {
441 c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
442 }
443}
444
445func TestSplitInBuckets(t *testing.T) {
446 c := qt.New(t)
447
448 nLeafs := 16
449 kvs := make([]kv, nLeafs)
450 for i := 0; i < nLeafs; i++ {
451 k := BigIntToBytes(big.NewInt(int64(i)))
452 v := BigIntToBytes(big.NewInt(int64(i * 2)))
453 keyPath := make([]byte, 32)
454 copy(keyPath[:], k)
455 kvs[i].pos = i
456 kvs[i].keyPath = k
457 kvs[i].k = k
458 kvs[i].v = v
459 }
460
461 // check keyToBucket results for 4 buckets & 8 keys
462 c.Assert(keyToBucket(kvs[0].k, 4), qt.Equals, 0)
463 c.Assert(keyToBucket(kvs[1].k, 4), qt.Equals, 2)
464 c.Assert(keyToBucket(kvs[2].k, 4), qt.Equals, 1)
465 c.Assert(keyToBucket(kvs[3].k, 4), qt.Equals, 3)
466 c.Assert(keyToBucket(kvs[4].k, 4), qt.Equals, 0)
467 c.Assert(keyToBucket(kvs[5].k, 4), qt.Equals, 2)
468 c.Assert(keyToBucket(kvs[6].k, 4), qt.Equals, 1)
469 c.Assert(keyToBucket(kvs[7].k, 4), qt.Equals, 3)
470
471 // check keyToBucket results for 8 buckets & 8 keys
472 c.Assert(keyToBucket(kvs[0].k, 8), qt.Equals, 0)
473 c.Assert(keyToBucket(kvs[1].k, 8), qt.Equals, 4)
474 c.Assert(keyToBucket(kvs[2].k, 8), qt.Equals, 2)
475 c.Assert(keyToBucket(kvs[3].k, 8), qt.Equals, 6)
476 c.Assert(keyToBucket(kvs[4].k, 8), qt.Equals, 1)
477 c.Assert(keyToBucket(kvs[5].k, 8), qt.Equals, 5)
478 c.Assert(keyToBucket(kvs[6].k, 8), qt.Equals, 3)
479 c.Assert(keyToBucket(kvs[7].k, 8), qt.Equals, 7)
480
481 buckets := splitInBuckets(kvs, 4)
482
483 expected := [][]string{
484 {
485 "00000000", // bucket 0
486 "08000000",
487 "04000000",
488 "0c000000",
489 },
490 {
491 "02000000", // bucket 1
492 "0a000000",
493 "06000000",
494 "0e000000",
495 },
496 {
497 "01000000", // bucket 2
498 "09000000",
499 "05000000",
500 "0d000000",
501 },
502 {
503 "03000000", // bucket 3
504 "0b000000",
505 "07000000",
506 "0f000000",
507 },
508 }
509
510 for i := 0; i < len(buckets); i++ {
511 sortKvs(buckets[i])
512 c.Assert(len(buckets[i]), qt.Equals, len(expected[i]))
513 for j := 0; j < len(buckets[i]); j++ {
514 c.Check(hex.EncodeToString(buckets[i][j].k[:4]), qt.Equals, expected[i][j])
515 }
516 }
517}
518
519func TestAddBatchCaseC(t *testing.T) {
520 c := qt.New(t)
521
522 nLeafs := 1024
523 initialNLeafs := 101 // TMP TODO use const minLeafsThreshold+1 once ready
524
525 tree1, tree2 := testInit(c, initialNLeafs)
526 tree2.dbgInit()
527
528 start := time.Now()
529 for i := initialNLeafs; i < nLeafs; i++ {
530 k := BigIntToBytes(big.NewInt(int64(i)))
531 v := BigIntToBytes(big.NewInt(int64(i * 2)))
532 if err := tree1.Add(k, v); err != nil {
533 t.Fatal(err)
534 }
535 }
536 time1 := time.Since(start)
537
538 // prepare the key-values to be added
539 var keys, values [][]byte
540 for i := initialNLeafs; i < nLeafs; i++ {
541 k := BigIntToBytes(big.NewInt(int64(i)))
542 v := BigIntToBytes(big.NewInt(int64(i * 2)))
543 keys = append(keys, k)
544 values = append(values, v)
545 }
546 start = time.Now()
547 indexes, err := tree2.AddBatch(keys, values)
548 c.Assert(err, qt.IsNil)
549 time2 := time.Since(start)
550 if debug {
551 debugTime("CASE C, AddBatch", time1, time2)
552 printTestContext(" ", nLeafs, "Poseidon", "memory")
553 tree2.dbg.print(" ")
554 }
555 c.Check(len(indexes), qt.Equals, 0)
556
557 // check that both trees roots are equal
558 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
559}
560
561func TestAddBatchCaseD(t *testing.T) {
562 c := qt.New(t)
563
564 nLeafs := 4096
565 initialNLeafs := 900
566
567 tree1, tree2 := testInit(c, initialNLeafs)
568 tree2.dbgInit()
569
570 start := time.Now()
571 for i := initialNLeafs; i < nLeafs; i++ {
572 k := BigIntToBytes(big.NewInt(int64(i)))
573 v := BigIntToBytes(big.NewInt(int64(i * 2)))
574 if err := tree1.Add(k, v); err != nil {
575 t.Fatal(err)
576 }
577 }
578 time1 := time.Since(start)
579
580 // prepare the key-values to be added
581 var keys, values [][]byte
582 for i := initialNLeafs; i < nLeafs; i++ {
583 k := BigIntToBytes(big.NewInt(int64(i)))
584 v := BigIntToBytes(big.NewInt(int64(i * 2)))
585 keys = append(keys, k)
586 values = append(values, v)
587 }
588 start = time.Now()
589 indexes, err := tree2.AddBatch(keys, values)
590 c.Assert(err, qt.IsNil)
591 time2 := time.Since(start)
592 if debug {
593 debugTime("CASE D, AddBatch", time1, time2)
594 printTestContext(" ", nLeafs, "Poseidon", "memory")
595 tree2.dbg.print(" ")
596 }
597 c.Check(len(indexes), qt.Equals, 0)
598
599 // check that both trees roots are equal
600 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
601}
602
603func TestAddBatchCaseE(t *testing.T) {
604 c := qt.New(t)
605
606 nLeafs := 4096
607 initialNLeafs := 900
608
609 tree1, _ := testInit(c, initialNLeafs)
610
611 start := time.Now()
612 for i := initialNLeafs; i < nLeafs; i++ {
613 k := BigIntToBytes(big.NewInt(int64(i)))
614 v := BigIntToBytes(big.NewInt(int64(i * 2)))
615 if err := tree1.Add(k, v); err != nil {
616 t.Fatal(err)
617 }
618 }
619 time1 := time.Since(start)
620
621 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
622 c.Assert(err, qt.IsNil)
623 defer tree2.db.Close()
624 tree2.dbgInit()
625
626 var keys, values [][]byte
627 // add the initial leafs to fill a bit the tree before calling the
628 // AddBatch method
629 for i := 0; i < initialNLeafs; i++ {
630 k := BigIntToBytes(big.NewInt(int64(i)))
631 v := BigIntToBytes(big.NewInt(int64(i * 2)))
632 // use only the keys of one bucket, store the not used ones for
633 // later
634 if i%4 != 0 {
635 keys = append(keys, k)
636 values = append(values, v)
637 continue
638 }
639 if err := tree2.Add(k, v); err != nil {
640 t.Fatal(err)
641 }
642 }
643
644 for i := initialNLeafs; i < nLeafs; i++ {
645 k := BigIntToBytes(big.NewInt(int64(i)))
646 v := BigIntToBytes(big.NewInt(int64(i * 2)))
647 keys = append(keys, k)
648 values = append(values, v)
649 }
650 start = time.Now()
651 indexes, err := tree2.AddBatch(keys, values)
652 c.Assert(err, qt.IsNil)
653 time2 := time.Since(start)
654 if debug {
655 debugTime("CASE E, AddBatch", time1, time2)
656 printTestContext(" ", nLeafs, "Poseidon", "memory")
657 tree2.dbg.print(" ")
658 }
659 c.Check(len(indexes), qt.Equals, 0)
660
661 // check that both trees roots are equal
662 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
663}
664
665func TestFlp2(t *testing.T) {
666 c := qt.New(t)
667 c.Assert(flp2(31), qt.Equals, 16)
668 c.Assert(flp2(32), qt.Equals, 32)
669 c.Assert(flp2(33), qt.Equals, 32)
670 c.Assert(flp2(63), qt.Equals, 32)
671 c.Assert(flp2(64), qt.Equals, 64)
672 c.Assert(flp2(9000), qt.Equals, 8192)
673}
674
675func TestAddBatchBench(t *testing.T) {
676 nLeafs := 50_000
677 printTestContext("TestAddBatchBench: ", nLeafs, "Blake2b", "leveldb")
678
679 // prepare inputs
680 var ks, vs [][]byte
681 for i := 0; i < nLeafs; i++ {
682 k := randomBytes(32)
683 v := randomBytes(32)
684 ks = append(ks, k)
685 vs = append(vs, v)
686 }
687
688 benchAdd(t, ks, vs)
689
690 benchAddBatch(t, ks, vs)
691}
692
693func benchAdd(t *testing.T, ks, vs [][]byte) {
694 c := qt.New(t)
695
696 dbDir := t.TempDir()
697 storage, err := leveldb.NewLevelDbStorage(dbDir, false)
698 c.Assert(err, qt.IsNil)
699 tree, err := NewTree(storage, 140, HashFunctionBlake2b)
700 c.Assert(err, qt.IsNil)
701
702 if debug {
703 tree.dbgInit()
704 }
705 start := time.Now()
706 for i := 0; i < len(ks); i++ {
707 err = tree.Add(ks[i], vs[i])
708 c.Assert(err, qt.IsNil)
709 }
710 printRes(" Add loop", time.Since(start))
711 tree.dbg.print(" ")
712}
713
714func benchAddBatch(t *testing.T, ks, vs [][]byte) {
715 c := qt.New(t)
716
717 dbDir := t.TempDir()
718 storage, err := leveldb.NewLevelDbStorage(dbDir, false)
719 c.Assert(err, qt.IsNil)
720 tree, err := NewTree(storage, 140, HashFunctionBlake2b)
721 c.Assert(err, qt.IsNil)
722
723 if debug {
724 tree.dbgInit()
725 }
726 start := time.Now()
727 invalids, err := tree.AddBatch(ks, vs)
728 printRes(" AddBatch", time.Since(start))
729 c.Assert(err, qt.IsNil)
730 c.Assert(len(invalids), qt.Equals, 0)
731 tree.dbg.print(" ")
732}
733
734func TestDbgStats(t *testing.T) {
735 c := qt.New(t)
736
737 nLeafs := 10_000
738
739 // prepare inputs
740 var ks, vs [][]byte
741 for i := 0; i < nLeafs; i++ {
742 k := randomBytes(32)
743 v := randomBytes(32)
744 ks = append(ks, k)
745 vs = append(vs, v)
746 }
747
748 // 1
749 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
750 c.Assert(err, qt.IsNil)
751 defer tree1.db.Close()
752
753 tree1.dbgInit()
754
755 for i := 0; i < len(ks); i++ {
756 err = tree1.Add(ks[i], vs[i])
757 c.Assert(err, qt.IsNil)
758 }
759
760 // 2
761 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
762 c.Assert(err, qt.IsNil)
763 defer tree2.db.Close()
764
765 tree2.dbgInit()
766
767 invalids, err := tree2.AddBatch(ks, vs)
768 c.Assert(err, qt.IsNil)
769 c.Assert(len(invalids), qt.Equals, 0)
770
771 // 3
772 tree3, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
773 c.Assert(err, qt.IsNil)
774 defer tree3.db.Close()
775
776 tree3.dbgInit()
777
778 // add few key-values
779 // invalids, err = tree3.AddBatch(ks[:], vs[:])
780 invalids, err = tree3.AddBatch(ks[:1000], vs[:1000])
781 c.Assert(err, qt.IsNil)
782 c.Assert(len(invalids), qt.Equals, 0)
783
784 // add the rest of key-values
785 invalids, err = tree3.AddBatch(ks[1000:], vs[1000:])
786 c.Assert(err, qt.IsNil)
787 c.Assert(len(invalids), qt.Equals, 0)
788
789 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
790 c.Check(tree3.Root(), qt.DeepEquals, tree1.Root())
791
792 if debug {
793 fmt.Println("TestDbgStats")
794 tree1.dbg.print(" add in loop ")
795 tree2.dbg.print(" addbatch caseA ")
796 tree3.dbg.print(" addbatch caseD ")
797 }
798}
799
800// func printLeafs(name string, t *Tree) {
801// w := bytes.NewBufferString("")
802//
803// err := t.Iterate(func(k, v []byte) {
804// if v[0] != PrefixValueLeaf {
805// return
806// }
807// leafK, _ := readLeafValue(v)
808// fmt.Fprintf(w, hex.EncodeToString(leafK[:4])+"\n")
809// })
810// if err != nil {
811// panic(err)
812// }
813// err = ioutil.WriteFile(name, w.Bytes(), 0644)
814// if err != nil {
815// panic(err)
816// }
817//
818// }
819
820// func TestComputeCosts(t *testing.T) {
821// fmt.Println(computeSimpleAddCost(10))
822// fmt.Println(computeFromLeafsAddCost(10))
823//
824// fmt.Println(computeSimpleAddCost(1024))
825// fmt.Println(computeFromLeafsAddCost(1024))
826// }
827
828// TODO test tree with nLeafs > minLeafsThreshold, but that at level L, there is
829// less keys than nBuckets (so CASE C could be applied if first few leafs are
830// added to balance the tree)
831
832// TODO test adding batch with repeated keys in the batch
833// TODO test adding batch with multiple invalid keys