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