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) {
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.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.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 TestAddBatchCaseB(t *testing.T) {
105 c := qt.New(t)
106
107 nLeafs := 1024
108
109 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
110 c.Assert(err, qt.IsNil)
111 defer tree.db.Close()
112
113 start := time.Now()
114 for i := 0; i < nLeafs; i++ {
115 k := BigIntToBytes(big.NewInt(int64(i)))
116 v := BigIntToBytes(big.NewInt(int64(i * 2)))
117 if err := tree.Add(k, v); err != nil {
118 t.Fatal(err)
119 }
120 }
121 fmt.Println(time.Since(start))
122
123 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
124 c.Assert(err, qt.IsNil)
125 defer tree2.db.Close()
126
127 // add the initial leafs to fill a bit the tree before calling the
128 // AddBatch method
129 for i := 0; i < 99; i++ { // TMP TODO use const minLeafsThreshold-1 once ready
130 k := BigIntToBytes(big.NewInt(int64(i)))
131 v := BigIntToBytes(big.NewInt(int64(i * 2)))
132 if err := tree2.Add(k, v); err != nil {
133 t.Fatal(err)
134 }
135 }
136
137 var keys, values [][]byte
138 for i := 99; i < nLeafs; i++ {
139 k := BigIntToBytes(big.NewInt(int64(i)))
140 v := BigIntToBytes(big.NewInt(int64(i * 2)))
141 keys = append(keys, k)
142 values = append(values, v)
143 }
144 start = time.Now()
145 indexes, err := tree2.AddBatchOpt(keys, values)
146 c.Assert(err, qt.IsNil)
147 fmt.Println(time.Since(start))
148 c.Check(len(indexes), qt.Equals, 0)
149
150 // check that both trees roots are equal
151 c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
152}
153
154func TestGetKeysAtLevel(t *testing.T) {
155 c := qt.New(t)
156
157 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
158 c.Assert(err, qt.IsNil)
159 defer tree.db.Close()
160
161 for i := 0; i < 32; i++ {
162 k := BigIntToBytes(big.NewInt(int64(i)))
163 v := BigIntToBytes(big.NewInt(int64(i * 2)))
164 if err := tree.Add(k, v); err != nil {
165 t.Fatal(err)
166 }
167 }
168
169 keys, err := tree.getKeysAtLevel(2)
170 c.Assert(err, qt.IsNil)
171 expected := []string{
172 "a5d5f14fce7348e40751496cf25d107d91b0bd043435b9577d778a01f8aa6111",
173 "e9e8dd9b28a7f81d1ff34cb5cefc0146dd848b31031a427b79bdadb62e7f6910",
174 }
175 for i := 0; i < len(keys); i++ {
176 c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
177 }
178
179 keys, err = tree.getKeysAtLevel(3)
180 c.Assert(err, qt.IsNil)
181 expected = []string{
182 "9f12c13e52bca96ad4882a26558e48ab67ddd63e062b839207e893d961390f01",
183 "16d246dd6826ec7346c7328f11c4261facf82d4689f33263ff6e207956a77f21",
184 "4a22cc901c6337daa17a431fa20170684b710e5f551509511492ec24e81a8f2f",
185 "470d61abcbd154977bffc9a9ec5a8daff0caabcf2a25e8441f604c79daa0f82d",
186 }
187 for i := 0; i < len(keys); i++ {
188 c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
189 }
190
191 keys, err = tree.getKeysAtLevel(4)
192 c.Assert(err, qt.IsNil)
193 expected = []string{
194 "7a5d1c81f7b96318012de3417e53d4f13df5b1337718651cd29d0cb0a66edd20",
195 "3408213e4e844bdf3355eb8781c74e31626812898c2dbe141ed6d2c92256fc1c",
196 "dfd8a4d0b6954a3e9f3892e655b58d456eeedf9367f27dfdd9bc2dd6a5577312",
197 "9e99fbec06fb2a6725997c12c4995f62725eb4cce4808523a5a5e80cca64b007",
198 "0befa1e070231dbf4e8ff841c05878cdec823e0c09594c24910a248b3ff5a628",
199 "b7131b0a15c772a57005a4dc5d0d6dd4b3414f5d9ee7408ce5e86c5ab3520e04",
200 "6d1abe0364077846a56bab1deb1a04883eb796b74fe531a7676a9a370f83ab21",
201 "4270116394bede69cf9cd72069eca018238080380bef5de75be8dcbbe968e105",
202 }
203 for i := 0; i < len(keys); i++ {
204 c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
205 }
206}
207
208func TestSplitInBuckets(t *testing.T) {
209 c := qt.New(t)
210
211 nLeafs := 16
212 kvs := make([]kv, nLeafs)
213 for i := 0; i < nLeafs; i++ {
214 k := BigIntToBytes(big.NewInt(int64(i)))
215 v := BigIntToBytes(big.NewInt(int64(i * 2)))
216 keyPath := make([]byte, 32)
217 copy(keyPath[:], k)
218 kvs[i].pos = i
219 kvs[i].keyPath = k
220 kvs[i].k = k
221 kvs[i].v = v
222 }
223
224 // check keyToBucket results for 4 buckets & 8 keys
225 c.Assert(keyToBucket(kvs[0].k, 4), qt.Equals, 0)
226 c.Assert(keyToBucket(kvs[1].k, 4), qt.Equals, 2)
227 c.Assert(keyToBucket(kvs[2].k, 4), qt.Equals, 1)
228 c.Assert(keyToBucket(kvs[3].k, 4), qt.Equals, 3)
229 c.Assert(keyToBucket(kvs[4].k, 4), qt.Equals, 0)
230 c.Assert(keyToBucket(kvs[5].k, 4), qt.Equals, 2)
231 c.Assert(keyToBucket(kvs[6].k, 4), qt.Equals, 1)
232 c.Assert(keyToBucket(kvs[7].k, 4), qt.Equals, 3)
233
234 // check keyToBucket results for 8 buckets & 8 keys
235 c.Assert(keyToBucket(kvs[0].k, 8), qt.Equals, 0)
236 c.Assert(keyToBucket(kvs[1].k, 8), qt.Equals, 4)
237 c.Assert(keyToBucket(kvs[2].k, 8), qt.Equals, 2)
238 c.Assert(keyToBucket(kvs[3].k, 8), qt.Equals, 6)
239 c.Assert(keyToBucket(kvs[4].k, 8), qt.Equals, 1)
240 c.Assert(keyToBucket(kvs[5].k, 8), qt.Equals, 5)
241 c.Assert(keyToBucket(kvs[6].k, 8), qt.Equals, 3)
242 c.Assert(keyToBucket(kvs[7].k, 8), qt.Equals, 7)
243
244 buckets := splitInBuckets(kvs, 4)
245
246 expected := [][]string{
247 {
248 "00000000", // bucket 0
249 "08000000",
250 "04000000",
251 "0c000000",
252 },
253 {
254 "02000000", // bucket 1
255 "0a000000",
256 "06000000",
257 "0e000000",
258 },
259 {
260 "01000000", // bucket 2
261 "09000000",
262 "05000000",
263 "0d000000",
264 },
265 {
266 "03000000", // bucket 3
267 "0b000000",
268 "07000000",
269 "0f000000",
270 },
271 }
272
273 for i := 0; i < len(buckets); i++ {
274 sortKvs(buckets[i])
275 c.Assert(len(buckets[i]), qt.Equals, len(expected[i]))
276 for j := 0; j < len(buckets[i]); j++ {
277 c.Check(hex.EncodeToString(buckets[i][j].k[:4]), qt.Equals, expected[i][j])
278 }
279 }
280}
281
282func TestAddBatchCaseC(t *testing.T) {
283 c := qt.New(t)
284
285 nLeafs := 1024
286
287 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
288 c.Assert(err, qt.IsNil)
289 defer tree.db.Close()
290
291 start := time.Now()
292 for i := 0; i < nLeafs; i++ {
293 k := BigIntToBytes(big.NewInt(int64(i)))
294 v := BigIntToBytes(big.NewInt(int64(i * 2)))
295 if err := tree.Add(k, v); err != nil {
296 t.Fatal(err)
297 }
298 }
299 fmt.Println(time.Since(start))
300
301 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
302 c.Assert(err, qt.IsNil)
303 defer tree2.db.Close()
304
305 // add the initial leafs to fill a bit the tree before calling the
306 // AddBatch method
307 for i := 0; i < 101; i++ { // TMP TODO use const minLeafsThreshold-1 once ready
308 k := BigIntToBytes(big.NewInt(int64(i)))
309 v := BigIntToBytes(big.NewInt(int64(i * 2)))
310 if err := tree2.Add(k, v); err != nil {
311 t.Fatal(err)
312 }
313 }
314 // tree2.PrintGraphviz(nil)
315
316 var keys, values [][]byte
317 for i := 101; i < nLeafs; i++ {
318 k := BigIntToBytes(big.NewInt(int64(i)))
319 v := BigIntToBytes(big.NewInt(int64(i * 2)))
320 keys = append(keys, k)
321 values = append(values, v)
322 }
323 start = time.Now()
324 indexes, err := tree2.AddBatchOpt(keys, values)
325 c.Assert(err, qt.IsNil)
326 fmt.Println(time.Since(start))
327 c.Check(len(indexes), qt.Equals, 0)
328
329 // check that both trees roots are equal
330 c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
331
332 // tree.PrintGraphviz(nil)
333 // tree2.PrintGraphviz(nil)
334 // // tree.PrintGraphvizFirstNLevels(nil, 4)
335 // // tree2.PrintGraphvizFirstNLevels(nil, 4)
336 // fmt.Println("TREE")
337 // printLeafs("t1.txt", tree)
338 // fmt.Println("TREE2")
339 // printLeafs("t2.txt", tree2)
340}
341
342func TestAddBatchCaseD(t *testing.T) {
343 c := qt.New(t)
344
345 nLeafs := 8192
346
347 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
348 c.Assert(err, qt.IsNil)
349 defer tree.db.Close()
350
351 start := time.Now()
352 for i := 0; i < nLeafs; i++ {
353 k := BigIntToBytes(big.NewInt(int64(i)))
354 v := BigIntToBytes(big.NewInt(int64(i * 2)))
355 if err := tree.Add(k, v); err != nil {
356 t.Fatal(err)
357 }
358 }
359 fmt.Println(time.Since(start))
360
361 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
362 c.Assert(err, qt.IsNil)
363 defer tree2.db.Close()
364
365 // add the initial leafs to fill a bit the tree before calling the
366 // AddBatch method
367 for i := 0; i < 900; i++ { // TMP TODO use const minLeafsThreshold-1 once ready
368 k := BigIntToBytes(big.NewInt(int64(i)))
369 v := BigIntToBytes(big.NewInt(int64(i * 2)))
370 if err := tree2.Add(k, v); err != nil {
371 t.Fatal(err)
372 }
373 }
374 // tree2.PrintGraphvizFirstNLevels(nil, 4)
375 // tree2.PrintGraphviz(nil)
376
377 var keys, values [][]byte
378 for i := 900; i < nLeafs; i++ {
379 k := BigIntToBytes(big.NewInt(int64(i)))
380 v := BigIntToBytes(big.NewInt(int64(i * 2)))
381 keys = append(keys, k)
382 values = append(values, v)
383 }
384 start = time.Now()
385 indexes, err := tree2.AddBatchOpt(keys, values)
386 c.Assert(err, qt.IsNil)
387 fmt.Println(time.Since(start))
388 c.Check(len(indexes), qt.Equals, 0)
389
390 // check that both trees roots are equal
391 c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
392}
393
394// func printLeafs(name string, t *Tree) {
395// w := bytes.NewBufferString("")
396//
397// err := t.Iterate(func(k, v []byte) {
398// if v[0] != PrefixValueLeaf {
399// return
400// }
401// leafK, _ := readLeafValue(v)
402// fmt.Fprintf(w, hex.EncodeToString(leafK[:4])+"\n")
403// })
404// if err != nil {
405// panic(err)
406// }
407// err = ioutil.WriteFile(name, w.Bytes(), 0644)
408// if err != nil {
409// panic(err)
410// }
411//
412// }
413
414// func TestComputeCosts(t *testing.T) {
415// fmt.Println(computeSimpleAddCost(10))
416// fmt.Println(computeBottomUpAddCost(10))
417//
418// fmt.Println(computeSimpleAddCost(1024))
419// fmt.Println(computeBottomUpAddCost(1024))
420// }
421
422// TODO test tree with nLeafs > minLeafsThreshold, but that at level L, there is
423// less keys than nBuckets (so CASE C could be applied if first few leafs are
424// added to balance the tree)