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 = true
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 // tree1.PrintGraphviz(nil)
234
235 indexes, err := tree2.AddBatch(keys, values)
236 c.Assert(err, qt.IsNil)
237 // tree1.PrintGraphviz(nil)
238 // tree2.PrintGraphviz(nil)
239
240 c.Check(len(indexes), qt.Equals, 0)
241
242 // tree1.PrintGraphviz(nil)
243 // tree2.PrintGraphviz(nil)
244
245 // check that both trees roots are equal
246 // fmt.Println(hex.EncodeToString(tree1.Root()))
247 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
248 // c.Assert(tree2.Root(), qt.DeepEquals, tree1.Root())
249
250 // fmt.Println("\n---2nd test vector---")
251 //
252 // // 2nd test vectors
253 // tree1, err = NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
254 // c.Assert(err, qt.IsNil)
255 // defer tree1.db.Close()
256 //
257 // tree2, err = NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
258 // c.Assert(err, qt.IsNil)
259 // defer tree2.db.Close()
260 //
261 // testvectorKeys = []string{
262 // "1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642",
263 // "2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf",
264 // "9cb87ec67e875c61390edcd1ab517f443591047709a4d4e45b0f9ed980857b8e",
265 // "9b4e9e92e974a589f426ceeb4cb291dc24893513fecf8e8460992dcf52621d4d",
266 // "1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5",
267 // "d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7",
268 // "3cd55dbfb8f975f20a0925dfbdabe79fa2d51dd0268afbb8ba6b01de9dfcdd3c",
269 // "5d0a9d6d9f197c091bf054fac9cb60e11ec723d6610ed8578e617b4d46cb43d5",
270 // }
271 // keys = [][]byte{}
272 // values = [][]byte{}
273 // for i := 0; i < len(testvectorKeys); i++ {
274 // key, err := hex.DecodeString(testvectorKeys[i])
275 // c.Assert(err, qt.IsNil)
276 // keys = append(keys, key)
277 // values = append(values, []byte{0})
278 // }
279 //
280 // for i := 0; i < len(keys); i++ {
281 // if err := tree1.Add(keys[i], values[i]); err != nil {
282 // t.Fatal(err)
283 // }
284 // }
285 //
286 // indexes, err = tree2.AddBatch(keys, values)
287 // c.Assert(err, qt.IsNil)
288 // // tree1.PrintGraphviz(nil)
289 // // tree2.PrintGraphviz(nil)
290 //
291 // c.Check(len(indexes), qt.Equals, 0)
292 //
293 // // check that both trees roots are equal
294 // // fmt.Println(hex.EncodeToString(tree1.Root()))
295 // c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
296}
297
298func TestAddBatchCaseARandomKeys(t *testing.T) {
299 c := qt.New(t)
300
301 nLeafs := 8
302
303 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
304 c.Assert(err, qt.IsNil)
305 defer tree1.db.Close()
306
307 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
308 c.Assert(err, qt.IsNil)
309 defer tree2.db.Close()
310
311 var keys, values [][]byte
312 for i := 0; i < nLeafs; i++ {
313 keys = append(keys, randomBytes(32))
314 // values = append(values, randomBytes(32))
315 values = append(values, []byte{0})
316 // fmt.Println("K", hex.EncodeToString(keys[i]))
317 }
318
319 // TODO delete
320 keys[0], _ = hex.DecodeString("1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642")
321 keys[1], _ = hex.DecodeString("2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf")
322 keys[2], _ = hex.DecodeString("9cb87ec67e875c61390edcd1ab517f443591047709a4d4e45b0f9ed980857b8e")
323 keys[3], _ = hex.DecodeString("9b4e9e92e974a589f426ceeb4cb291dc24893513fecf8e8460992dcf52621d4d")
324 keys[4], _ = hex.DecodeString("1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5")
325 keys[5], _ = hex.DecodeString("d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7")
326 keys[6], _ = hex.DecodeString("3cd55dbfb8f975f20a0925dfbdabe79fa2d51dd0268afbb8ba6b01de9dfcdd3c")
327 keys[7], _ = hex.DecodeString("5d0a9d6d9f197c091bf054fac9cb60e11ec723d6610ed8578e617b4d46cb43d5")
328
329 for i := 0; i < len(keys); i++ {
330 if err := tree1.Add(keys[i], values[i]); err != nil {
331 t.Fatal(err)
332 }
333 }
334
335 indexes, err := tree2.AddBatch(keys, values)
336 c.Assert(err, qt.IsNil)
337 // tree1.PrintGraphviz(nil)
338 // tree2.PrintGraphviz(nil)
339
340 c.Check(len(indexes), qt.Equals, 0)
341
342 // check that both trees roots are equal
343 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
344}
345
346func TestAddBatchCaseB(t *testing.T) {
347 c := qt.New(t)
348
349 nLeafs := 1024
350 initialNLeafs := 99 // TMP TODO use const minLeafsThreshold-1 once ready
351
352 tree1, tree2 := testInit(c, initialNLeafs)
353 tree2.dbgInit()
354
355 start := time.Now()
356 for i := initialNLeafs; i < nLeafs; i++ {
357 k := BigIntToBytes(big.NewInt(int64(i)))
358 v := BigIntToBytes(big.NewInt(int64(i * 2)))
359 if err := tree1.Add(k, v); err != nil {
360 t.Fatal(err)
361 }
362 }
363 time1 := time.Since(start)
364
365 // prepare the key-values to be added
366 var keys, values [][]byte
367 for i := initialNLeafs; i < nLeafs; i++ {
368 k := BigIntToBytes(big.NewInt(int64(i)))
369 v := BigIntToBytes(big.NewInt(int64(i * 2)))
370 keys = append(keys, k)
371 values = append(values, v)
372 }
373 start = time.Now()
374 indexes, err := tree2.AddBatch(keys, values)
375 c.Assert(err, qt.IsNil)
376 time2 := time.Since(start)
377 if debug {
378 debugTime("CASE B, AddBatch", time1, time2)
379 printTestContext(" ", nLeafs, "Poseidon", "memory")
380 tree2.dbg.print(" ")
381 }
382 c.Check(len(indexes), qt.Equals, 0)
383
384 // check that both trees roots are equal
385 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
386}
387
388func TestAddBatchCaseBRepeatedLeafs(t *testing.T) {
389 c := qt.New(t)
390
391 nLeafs := 1024
392 initialNLeafs := 99 // TMP TODO use const minLeafsThreshold-1 once ready
393
394 tree1, tree2 := testInit(c, initialNLeafs)
395
396 for i := initialNLeafs; i < nLeafs; i++ {
397 k := BigIntToBytes(big.NewInt(int64(i)))
398 v := BigIntToBytes(big.NewInt(int64(i * 2)))
399 if err := tree1.Add(k, v); err != nil {
400 t.Fatal(err)
401 }
402 }
403
404 // prepare the key-values to be added, including already existing keys
405 var keys, values [][]byte
406 for i := 0; i < nLeafs; i++ {
407 k := BigIntToBytes(big.NewInt(int64(i)))
408 v := BigIntToBytes(big.NewInt(int64(i * 2)))
409 keys = append(keys, k)
410 values = append(values, v)
411 }
412 indexes, err := tree2.AddBatch(keys, values)
413 c.Assert(err, qt.IsNil)
414 c.Check(len(indexes), qt.Equals, initialNLeafs)
415
416 // check that both trees roots are equal
417 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
418}
419
420func TestCombineInKVSet(t *testing.T) {
421 c := qt.New(t)
422
423 var a, b, expected []kv
424 for i := 0; i < 10; i++ {
425 k := BigIntToBytes(big.NewInt(int64(i)))
426 kv := kv{k: k}
427 if i < 7 {
428 a = append(a, kv)
429 }
430 if i >= 4 {
431 b = append(b, kv)
432 }
433 expected = append(expected, kv)
434 }
435
436 r, invalids := combineInKVSet(a, b)
437 for i := 0; i < len(r); i++ {
438 c.Assert(r[i].k, qt.DeepEquals, expected[i].k)
439 }
440 c.Assert(len(invalids), qt.Equals, 7-4)
441}
442
443func TestGetKeysAtLevel(t *testing.T) {
444 c := qt.New(t)
445
446 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
447 c.Assert(err, qt.IsNil)
448 defer tree.db.Close()
449
450 for i := 0; i < 32; i++ {
451 k := BigIntToBytes(big.NewInt(int64(i)))
452 v := BigIntToBytes(big.NewInt(int64(i * 2)))
453 if err := tree.Add(k, v); err != nil {
454 t.Fatal(err)
455 }
456 }
457
458 keys, err := tree.getKeysAtLevel(2)
459 c.Assert(err, qt.IsNil)
460 expected := []string{
461 "a5d5f14fce7348e40751496cf25d107d91b0bd043435b9577d778a01f8aa6111",
462 "e9e8dd9b28a7f81d1ff34cb5cefc0146dd848b31031a427b79bdadb62e7f6910",
463 }
464 for i := 0; i < len(keys); i++ {
465 c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
466 }
467
468 keys, err = tree.getKeysAtLevel(3)
469 c.Assert(err, qt.IsNil)
470 expected = []string{
471 "9f12c13e52bca96ad4882a26558e48ab67ddd63e062b839207e893d961390f01",
472 "16d246dd6826ec7346c7328f11c4261facf82d4689f33263ff6e207956a77f21",
473 "4a22cc901c6337daa17a431fa20170684b710e5f551509511492ec24e81a8f2f",
474 "470d61abcbd154977bffc9a9ec5a8daff0caabcf2a25e8441f604c79daa0f82d",
475 }
476 for i := 0; i < len(keys); i++ {
477 c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
478 }
479
480 keys, err = tree.getKeysAtLevel(4)
481 c.Assert(err, qt.IsNil)
482 expected = []string{
483 "7a5d1c81f7b96318012de3417e53d4f13df5b1337718651cd29d0cb0a66edd20",
484 "3408213e4e844bdf3355eb8781c74e31626812898c2dbe141ed6d2c92256fc1c",
485 "dfd8a4d0b6954a3e9f3892e655b58d456eeedf9367f27dfdd9bc2dd6a5577312",
486 "9e99fbec06fb2a6725997c12c4995f62725eb4cce4808523a5a5e80cca64b007",
487 "0befa1e070231dbf4e8ff841c05878cdec823e0c09594c24910a248b3ff5a628",
488 "b7131b0a15c772a57005a4dc5d0d6dd4b3414f5d9ee7408ce5e86c5ab3520e04",
489 "6d1abe0364077846a56bab1deb1a04883eb796b74fe531a7676a9a370f83ab21",
490 "4270116394bede69cf9cd72069eca018238080380bef5de75be8dcbbe968e105",
491 }
492 for i := 0; i < len(keys); i++ {
493 c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
494 }
495}
496
497func TestSplitInBuckets(t *testing.T) {
498 c := qt.New(t)
499
500 nLeafs := 16
501 kvs := make([]kv, nLeafs)
502 for i := 0; i < nLeafs; i++ {
503 k := BigIntToBytes(big.NewInt(int64(i)))
504 v := BigIntToBytes(big.NewInt(int64(i * 2)))
505 keyPath := make([]byte, 32)
506 copy(keyPath[:], k)
507 kvs[i].pos = i
508 kvs[i].keyPath = k
509 kvs[i].k = k
510 kvs[i].v = v
511 }
512
513 // check keyToBucket results for 4 buckets & 8 keys
514 c.Assert(keyToBucket(kvs[0].k, 4), qt.Equals, 0)
515 c.Assert(keyToBucket(kvs[1].k, 4), qt.Equals, 2)
516 c.Assert(keyToBucket(kvs[2].k, 4), qt.Equals, 1)
517 c.Assert(keyToBucket(kvs[3].k, 4), qt.Equals, 3)
518 c.Assert(keyToBucket(kvs[4].k, 4), qt.Equals, 0)
519 c.Assert(keyToBucket(kvs[5].k, 4), qt.Equals, 2)
520 c.Assert(keyToBucket(kvs[6].k, 4), qt.Equals, 1)
521 c.Assert(keyToBucket(kvs[7].k, 4), qt.Equals, 3)
522
523 // check keyToBucket results for 8 buckets & 8 keys
524 c.Assert(keyToBucket(kvs[0].k, 8), qt.Equals, 0)
525 c.Assert(keyToBucket(kvs[1].k, 8), qt.Equals, 4)
526 c.Assert(keyToBucket(kvs[2].k, 8), qt.Equals, 2)
527 c.Assert(keyToBucket(kvs[3].k, 8), qt.Equals, 6)
528 c.Assert(keyToBucket(kvs[4].k, 8), qt.Equals, 1)
529 c.Assert(keyToBucket(kvs[5].k, 8), qt.Equals, 5)
530 c.Assert(keyToBucket(kvs[6].k, 8), qt.Equals, 3)
531 c.Assert(keyToBucket(kvs[7].k, 8), qt.Equals, 7)
532
533 buckets := splitInBuckets(kvs, 4)
534
535 expected := [][]string{
536 {
537 "00000000", // bucket 0
538 "08000000",
539 "04000000",
540 "0c000000",
541 },
542 {
543 "02000000", // bucket 1
544 "0a000000",
545 "06000000",
546 "0e000000",
547 },
548 {
549 "01000000", // bucket 2
550 "09000000",
551 "05000000",
552 "0d000000",
553 },
554 {
555 "03000000", // bucket 3
556 "0b000000",
557 "07000000",
558 "0f000000",
559 },
560 }
561
562 for i := 0; i < len(buckets); i++ {
563 sortKvs(buckets[i])
564 c.Assert(len(buckets[i]), qt.Equals, len(expected[i]))
565 for j := 0; j < len(buckets[i]); j++ {
566 c.Check(hex.EncodeToString(buckets[i][j].k[:4]), qt.Equals, expected[i][j])
567 }
568 }
569}
570
571func TestAddBatchCaseC(t *testing.T) {
572 c := qt.New(t)
573
574 nLeafs := 1024
575 initialNLeafs := 101 // TMP TODO use const minLeafsThreshold+1 once ready
576
577 tree1, tree2 := testInit(c, initialNLeafs)
578 tree2.dbgInit()
579
580 start := time.Now()
581 for i := initialNLeafs; i < nLeafs; i++ {
582 k := BigIntToBytes(big.NewInt(int64(i)))
583 v := BigIntToBytes(big.NewInt(int64(i * 2)))
584 if err := tree1.Add(k, v); err != nil {
585 t.Fatal(err)
586 }
587 }
588 time1 := time.Since(start)
589
590 // prepare the key-values to be added
591 var keys, values [][]byte
592 for i := initialNLeafs; i < nLeafs; i++ {
593 k := BigIntToBytes(big.NewInt(int64(i)))
594 v := BigIntToBytes(big.NewInt(int64(i * 2)))
595 keys = append(keys, k)
596 values = append(values, v)
597 }
598 start = time.Now()
599 indexes, err := tree2.AddBatch(keys, values)
600 c.Assert(err, qt.IsNil)
601 time2 := time.Since(start)
602 if debug {
603 debugTime("CASE C, AddBatch", time1, time2)
604 printTestContext(" ", nLeafs, "Poseidon", "memory")
605 tree2.dbg.print(" ")
606 }
607 c.Check(len(indexes), qt.Equals, 0)
608
609 // check that both trees roots are equal
610 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
611}
612
613func TestAddBatchCaseD(t *testing.T) {
614 c := qt.New(t)
615
616 nLeafs := 4096
617 initialNLeafs := 900
618
619 tree1, tree2 := testInit(c, initialNLeafs)
620 tree2.dbgInit()
621
622 start := time.Now()
623 for i := initialNLeafs; i < nLeafs; i++ {
624 k := BigIntToBytes(big.NewInt(int64(i)))
625 v := BigIntToBytes(big.NewInt(int64(i * 2)))
626 if err := tree1.Add(k, v); err != nil {
627 t.Fatal(err)
628 }
629 }
630 time1 := time.Since(start)
631
632 // prepare the key-values to be added
633 var keys, values [][]byte
634 for i := initialNLeafs; i < nLeafs; i++ {
635 k := BigIntToBytes(big.NewInt(int64(i)))
636 v := BigIntToBytes(big.NewInt(int64(i * 2)))
637 keys = append(keys, k)
638 values = append(values, v)
639 }
640 start = time.Now()
641 indexes, err := tree2.AddBatch(keys, values)
642 c.Assert(err, qt.IsNil)
643 time2 := time.Since(start)
644 if debug {
645 debugTime("CASE D, AddBatch", time1, time2)
646 printTestContext(" ", nLeafs, "Poseidon", "memory")
647 tree2.dbg.print(" ")
648 }
649 c.Check(len(indexes), qt.Equals, 0)
650
651 // check that both trees roots are equal
652 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
653}
654
655func TestAddBatchCaseE(t *testing.T) {
656 c := qt.New(t)
657
658 nLeafs := 4096
659 initialNLeafs := 900
660
661 tree1, _ := testInit(c, initialNLeafs)
662
663 start := time.Now()
664 for i := initialNLeafs; i < nLeafs; i++ {
665 k := BigIntToBytes(big.NewInt(int64(i)))
666 v := BigIntToBytes(big.NewInt(int64(i * 2)))
667 if err := tree1.Add(k, v); err != nil {
668 t.Fatal(err)
669 }
670 }
671 time1 := time.Since(start)
672
673 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
674 c.Assert(err, qt.IsNil)
675 defer tree2.db.Close()
676 tree2.dbgInit()
677
678 var keys, values [][]byte
679 // add the initial leafs to fill a bit the tree before calling the
680 // AddBatch method
681 for i := 0; i < initialNLeafs; i++ {
682 k := BigIntToBytes(big.NewInt(int64(i)))
683 v := BigIntToBytes(big.NewInt(int64(i * 2)))
684 // use only the keys of one bucket, store the not used ones for
685 // later
686 if i%4 != 0 {
687 keys = append(keys, k)
688 values = append(values, v)
689 continue
690 }
691 if err := tree2.Add(k, v); err != nil {
692 t.Fatal(err)
693 }
694 }
695
696 for i := initialNLeafs; i < nLeafs; i++ {
697 k := BigIntToBytes(big.NewInt(int64(i)))
698 v := BigIntToBytes(big.NewInt(int64(i * 2)))
699 keys = append(keys, k)
700 values = append(values, v)
701 }
702 start = time.Now()
703 indexes, err := tree2.AddBatch(keys, values)
704 c.Assert(err, qt.IsNil)
705 time2 := time.Since(start)
706 if debug {
707 debugTime("CASE E, AddBatch", time1, time2)
708 printTestContext(" ", nLeafs, "Poseidon", "memory")
709 tree2.dbg.print(" ")
710 }
711 c.Check(len(indexes), qt.Equals, 0)
712
713 // check that both trees roots are equal
714 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
715}
716
717func TestFlp2(t *testing.T) {
718 c := qt.New(t)
719 c.Assert(flp2(31), qt.Equals, 16)
720 c.Assert(flp2(32), qt.Equals, 32)
721 c.Assert(flp2(33), qt.Equals, 32)
722 c.Assert(flp2(63), qt.Equals, 32)
723 c.Assert(flp2(64), qt.Equals, 64)
724 c.Assert(flp2(9000), qt.Equals, 8192)
725}
726
727func TestAddBatchBench(t *testing.T) {
728 nLeafs := 50_000
729 printTestContext("TestAddBatchBench: ", nLeafs, "Blake2b", "leveldb")
730
731 // prepare inputs
732 var ks, vs [][]byte
733 for i := 0; i < nLeafs; i++ {
734 k := randomBytes(32)
735 v := randomBytes(32)
736 ks = append(ks, k)
737 vs = append(vs, v)
738 }
739
740 benchAdd(t, ks, vs)
741
742 benchAddBatch(t, ks, vs)
743}
744
745func benchAdd(t *testing.T, ks, vs [][]byte) {
746 c := qt.New(t)
747
748 dbDir := t.TempDir()
749 storage, err := leveldb.NewLevelDbStorage(dbDir, false)
750 c.Assert(err, qt.IsNil)
751 tree, err := NewTree(storage, 140, HashFunctionBlake2b)
752 c.Assert(err, qt.IsNil)
753
754 if debug {
755 tree.dbgInit()
756 }
757 start := time.Now()
758 for i := 0; i < len(ks); i++ {
759 err = tree.Add(ks[i], vs[i])
760 c.Assert(err, qt.IsNil)
761 }
762 printRes(" Add loop", time.Since(start))
763 tree.dbg.print(" ")
764}
765
766func benchAddBatch(t *testing.T, ks, vs [][]byte) {
767 c := qt.New(t)
768
769 dbDir := t.TempDir()
770 storage, err := leveldb.NewLevelDbStorage(dbDir, false)
771 c.Assert(err, qt.IsNil)
772 tree, err := NewTree(storage, 140, HashFunctionBlake2b)
773 c.Assert(err, qt.IsNil)
774
775 if debug {
776 tree.dbgInit()
777 }
778 start := time.Now()
779 invalids, err := tree.AddBatch(ks, vs)
780 printRes(" AddBatch", time.Since(start))
781 c.Assert(err, qt.IsNil)
782 c.Assert(len(invalids), qt.Equals, 0)
783 tree.dbg.print(" ")
784}
785
786func TestDbgStats(t *testing.T) {
787 c := qt.New(t)
788
789 nLeafs := 10_000
790
791 // prepare inputs
792 var ks, vs [][]byte
793 for i := 0; i < nLeafs; i++ {
794 k := randomBytes(32)
795 v := randomBytes(32)
796 ks = append(ks, k)
797 vs = append(vs, v)
798 }
799
800 // 1
801 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
802 c.Assert(err, qt.IsNil)
803 defer tree1.db.Close()
804
805 tree1.dbgInit()
806
807 for i := 0; i < len(ks); i++ {
808 err = tree1.Add(ks[i], vs[i])
809 c.Assert(err, qt.IsNil)
810 }
811
812 // 2
813 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
814 c.Assert(err, qt.IsNil)
815 defer tree2.db.Close()
816
817 tree2.dbgInit()
818
819 invalids, err := tree2.AddBatch(ks, vs)
820 c.Assert(err, qt.IsNil)
821 c.Assert(len(invalids), qt.Equals, 0)
822
823 // 3
824 tree3, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionBlake2b)
825 c.Assert(err, qt.IsNil)
826 defer tree3.db.Close()
827
828 tree3.dbgInit()
829
830 // add few key-values
831 // invalids, err = tree3.AddBatch(ks[:], vs[:])
832 invalids, err = tree3.AddBatch(ks[:1000], vs[:1000])
833 c.Assert(err, qt.IsNil)
834 c.Assert(len(invalids), qt.Equals, 0)
835
836 // add the rest of key-values
837 invalids, err = tree3.AddBatch(ks[1000:], vs[1000:])
838 c.Assert(err, qt.IsNil)
839 c.Assert(len(invalids), qt.Equals, 0)
840
841 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
842 c.Check(tree3.Root(), qt.DeepEquals, tree1.Root())
843
844 if debug {
845 fmt.Println("TestDbgStats")
846 tree1.dbg.print(" add in loop ")
847 tree2.dbg.print(" addbatch caseA ")
848 tree3.dbg.print(" addbatch caseD ")
849 }
850}
851
852func TestLoadVT(t *testing.T) {
853 c := qt.New(t)
854
855 nLeafs := 1024
856
857 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
858 c.Assert(err, qt.IsNil)
859 defer tree.db.Close()
860
861 var keys, values [][]byte
862 for i := 0; i < nLeafs; i++ {
863 k := randomBytes(31)
864 v := randomBytes(31)
865 keys = append(keys, k)
866 values = append(values, v)
867 }
868 indexes, err := tree.AddBatch(keys, values)
869 c.Assert(err, qt.IsNil)
870 c.Check(len(indexes), qt.Equals, 0)
871
872 vt, err := tree.loadVT()
873 c.Assert(err, qt.IsNil)
874 _, err = vt.computeHashes()
875 c.Assert(err, qt.IsNil)
876
877 // check that tree & vt roots are equal
878 c.Check(tree.Root(), qt.DeepEquals, vt.root.h)
879}
880
881// func printLeafs(name string, t *Tree) {
882// w := bytes.NewBufferString("")
883//
884// err := t.Iterate(func(k, v []byte) {
885// if v[0] != PrefixValueLeaf {
886// return
887// }
888// leafK, _ := readLeafValue(v)
889// fmt.Fprintf(w, hex.EncodeToString(leafK[:4])+"\n")
890// })
891// if err != nil {
892// panic(err)
893// }
894// err = ioutil.WriteFile(name, w.Bytes(), 0644)
895// if err != nil {
896// panic(err)
897// }
898//
899// }
900
901// func TestComputeCosts(t *testing.T) {
902// fmt.Println(computeSimpleAddCost(10))
903// fmt.Println(computeFromLeafsAddCost(10))
904//
905// fmt.Println(computeSimpleAddCost(1024))
906// fmt.Println(computeFromLeafsAddCost(1024))
907// }
908
909// TODO test tree with nLeafs > minLeafsThreshold, but that at level L, there is
910// less keys than nBuckets (so CASE C could be applied if first few leafs are
911// added to balance the tree)
912
913// TODO test adding batch with repeated keys in the batch
914// TODO test adding batch with multiple invalid keys