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