arbo/tree_test.go
tree_test.goRaw
1package arbo
2
3import (
4 "encoding/hex"
5 "math/big"
6 "testing"
7 "time"
8
9 qt "github.com/frankban/quicktest"
10 "go.vocdoni.io/dvote/db"
11)
12
13func TestDBBatch(t *testing.T) {
14 c := qt.New(t)
15
16 database, err := db.NewBadgerDB(c.TempDir())
17 c.Assert(err, qt.IsNil)
18 b := database.NewBatch()
19 err = b.Put([]byte("a"), []byte("b"))
20 c.Assert(err, qt.IsNil)
21}
22
23func TestAddTestVectors(t *testing.T) {
24 c := qt.New(t)
25
26 // Poseidon test vectors generated using https://github.com/iden3/circomlib smt.js
27 testVectorsPoseidon := []string{
28 "0000000000000000000000000000000000000000000000000000000000000000",
29 "13578938674299138072471463694055224830892726234048532520316387704878000008795",
30 "5412393676474193513566895793055462193090331607895808993925969873307089394741",
31 "14204494359367183802864593755198662203838502594566452929175967972147978322084",
32 }
33 testAdd(c, HashFunctionPoseidon, testVectorsPoseidon)
34
35 testVectorsSha256 := []string{
36 "0000000000000000000000000000000000000000000000000000000000000000",
37 "46910109172468462938850740851377282682950237270676610513794735904325820156367",
38 "59481735341404520835410489183267411392292882901306595567679529387376287440550",
39 "20573794434149960984975763118181266662429997821552560184909083010514790081771",
40 }
41 testAdd(c, HashFunctionSha256, testVectorsSha256)
42}
43
44func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
45 database, err := db.NewBadgerDB(c.TempDir())
46 c.Assert(err, qt.IsNil)
47 tree, err := NewTree(database, 10, hashFunc)
48 c.Assert(err, qt.IsNil)
49 defer tree.db.Close() //nolint:errcheck
50
51 c.Check(hex.EncodeToString(tree.Root()), qt.Equals, testVectors[0])
52
53 bLen := hashFunc.Len()
54 err = tree.Add(
55 BigIntToBytes(bLen, big.NewInt(1)),
56 BigIntToBytes(bLen, big.NewInt(2)))
57 c.Assert(err, qt.IsNil)
58 rootBI := BytesToBigInt(tree.Root())
59 c.Check(rootBI.String(), qt.Equals, testVectors[1])
60
61 err = tree.Add(
62 BigIntToBytes(bLen, big.NewInt(33)),
63 BigIntToBytes(bLen, big.NewInt(44)))
64 c.Assert(err, qt.IsNil)
65 rootBI = BytesToBigInt(tree.Root())
66 c.Check(rootBI.String(), qt.Equals, testVectors[2])
67
68 err = tree.Add(
69 BigIntToBytes(bLen, big.NewInt(1234)),
70 BigIntToBytes(bLen, big.NewInt(9876)))
71 c.Assert(err, qt.IsNil)
72 rootBI = BytesToBigInt(tree.Root())
73 c.Check(rootBI.String(), qt.Equals, testVectors[3])
74}
75
76func TestAddBatch(t *testing.T) {
77 c := qt.New(t)
78 database, err := db.NewBadgerDB(c.TempDir())
79 c.Assert(err, qt.IsNil)
80 tree, err := NewTree(database, 100, HashFunctionPoseidon)
81 c.Assert(err, qt.IsNil)
82 defer tree.db.Close() //nolint:errcheck
83
84 bLen := tree.HashFunction().Len()
85 for i := 0; i < 1000; i++ {
86 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
87 v := BigIntToBytes(bLen, big.NewInt(0))
88 if err := tree.Add(k, v); err != nil {
89 t.Fatal(err)
90 }
91 }
92
93 rootBI := BytesToBigInt(tree.Root())
94 c.Check(rootBI.String(), qt.Equals,
95 "296519252211642170490407814696803112091039265640052570497930797516015811235")
96
97 database, err = db.NewBadgerDB(c.TempDir())
98 c.Assert(err, qt.IsNil)
99 tree2, err := NewTree(database, 100, HashFunctionPoseidon)
100 c.Assert(err, qt.IsNil)
101 defer tree2.db.Close() //nolint:errcheck
102
103 var keys, values [][]byte
104 for i := 0; i < 1000; i++ {
105 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
106 v := BigIntToBytes(bLen, big.NewInt(0))
107 keys = append(keys, k)
108 values = append(values, v)
109 }
110 indexes, err := tree2.AddBatch(keys, values)
111 c.Assert(err, qt.IsNil)
112 c.Check(len(indexes), qt.Equals, 0)
113
114 rootBI = BytesToBigInt(tree2.Root())
115 c.Check(rootBI.String(), qt.Equals,
116 "296519252211642170490407814696803112091039265640052570497930797516015811235")
117}
118
119func TestAddDifferentOrder(t *testing.T) {
120 c := qt.New(t)
121 database1, err := db.NewBadgerDB(c.TempDir())
122 c.Assert(err, qt.IsNil)
123 tree1, err := NewTree(database1, 100, HashFunctionPoseidon)
124 c.Assert(err, qt.IsNil)
125 defer tree1.db.Close() //nolint:errcheck
126
127 bLen := tree1.HashFunction().Len()
128 for i := 0; i < 16; i++ {
129 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
130 v := BigIntToBytes(bLen, big.NewInt(0))
131 if err := tree1.Add(k, v); err != nil {
132 t.Fatal(err)
133 }
134 }
135
136 database2, err := db.NewBadgerDB(c.TempDir())
137 c.Assert(err, qt.IsNil)
138 tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
139 c.Assert(err, qt.IsNil)
140 defer tree2.db.Close() //nolint:errcheck
141
142 for i := 16 - 1; i >= 0; i-- {
143 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
144 v := BigIntToBytes(bLen, big.NewInt(0))
145 if err := tree2.Add(k, v); err != nil {
146 t.Fatal(err)
147 }
148 }
149
150 c.Check(hex.EncodeToString(tree2.Root()), qt.Equals, hex.EncodeToString(tree1.Root()))
151 c.Check(hex.EncodeToString(tree1.Root()), qt.Equals,
152 "3b89100bec24da9275c87bc188740389e1d5accfc7d88ba5688d7fa96a00d82f")
153}
154
155func TestAddRepeatedIndex(t *testing.T) {
156 c := qt.New(t)
157 database, err := db.NewBadgerDB(c.TempDir())
158 c.Assert(err, qt.IsNil)
159 tree, err := NewTree(database, 100, HashFunctionPoseidon)
160 c.Assert(err, qt.IsNil)
161 defer tree.db.Close() //nolint:errcheck
162
163 bLen := tree.HashFunction().Len()
164 k := BigIntToBytes(bLen, big.NewInt(int64(3)))
165 v := BigIntToBytes(bLen, big.NewInt(int64(12)))
166
167 err = tree.Add(k, v)
168 c.Assert(err, qt.IsNil)
169 err = tree.Add(k, v) // repeating same key-value
170 c.Check(err, qt.Equals, ErrKeyAlreadyExists)
171}
172
173func TestUpdate(t *testing.T) {
174 c := qt.New(t)
175 database, err := db.NewBadgerDB(c.TempDir())
176 c.Assert(err, qt.IsNil)
177 tree, err := NewTree(database, 100, HashFunctionPoseidon)
178 c.Assert(err, qt.IsNil)
179 defer tree.db.Close() //nolint:errcheck
180
181 bLen := tree.HashFunction().Len()
182 k := BigIntToBytes(bLen, big.NewInt(int64(20)))
183 v := BigIntToBytes(bLen, big.NewInt(int64(12)))
184 if err := tree.Add(k, v); err != nil {
185 t.Fatal(err)
186 }
187
188 v = BigIntToBytes(bLen, big.NewInt(int64(11)))
189 err = tree.Update(k, v)
190 c.Assert(err, qt.IsNil)
191
192 gettedKey, gettedValue, err := tree.Get(k)
193 c.Assert(err, qt.IsNil)
194 c.Check(gettedKey, qt.DeepEquals, k)
195 c.Check(gettedValue, qt.DeepEquals, v)
196
197 // add more leafs to the tree to do another test
198 for i := 0; i < 16; i++ {
199 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
200 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
201 if err := tree.Add(k, v); err != nil {
202 t.Fatal(err)
203 }
204 }
205
206 k = BigIntToBytes(bLen, big.NewInt(int64(3)))
207 v = BigIntToBytes(bLen, big.NewInt(int64(11)))
208 // check that before the Update, value for 3 is !=11
209 gettedKey, gettedValue, err = tree.Get(k)
210 c.Assert(err, qt.IsNil)
211 c.Check(gettedKey, qt.DeepEquals, k)
212 c.Check(gettedValue, qt.Not(qt.DeepEquals), v)
213 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(6)))
214
215 err = tree.Update(k, v)
216 c.Assert(err, qt.IsNil)
217
218 // check that after Update, the value for 3 is ==11
219 gettedKey, gettedValue, err = tree.Get(k)
220 c.Assert(err, qt.IsNil)
221 c.Check(gettedKey, qt.DeepEquals, k)
222 c.Check(gettedValue, qt.DeepEquals, v)
223 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(11)))
224}
225
226func TestAux(t *testing.T) { // TODO split in proper tests
227 c := qt.New(t)
228 database, err := db.NewBadgerDB(c.TempDir())
229 c.Assert(err, qt.IsNil)
230 tree, err := NewTree(database, 100, HashFunctionPoseidon)
231 c.Assert(err, qt.IsNil)
232 defer tree.db.Close() //nolint:errcheck
233
234 bLen := tree.HashFunction().Len()
235 k := BigIntToBytes(bLen, big.NewInt(int64(1)))
236 v := BigIntToBytes(bLen, big.NewInt(int64(0)))
237 err = tree.Add(k, v)
238 c.Assert(err, qt.IsNil)
239 k = BigIntToBytes(bLen, big.NewInt(int64(256)))
240 err = tree.Add(k, v)
241 c.Assert(err, qt.IsNil)
242
243 k = BigIntToBytes(bLen, big.NewInt(int64(257)))
244 err = tree.Add(k, v)
245 c.Assert(err, qt.IsNil)
246
247 k = BigIntToBytes(bLen, big.NewInt(int64(515)))
248 err = tree.Add(k, v)
249 c.Assert(err, qt.IsNil)
250 k = BigIntToBytes(bLen, big.NewInt(int64(770)))
251 err = tree.Add(k, v)
252 c.Assert(err, qt.IsNil)
253
254 k = BigIntToBytes(bLen, big.NewInt(int64(388)))
255 err = tree.Add(k, v)
256 c.Assert(err, qt.IsNil)
257 k = BigIntToBytes(bLen, big.NewInt(int64(900)))
258 err = tree.Add(k, v)
259 c.Assert(err, qt.IsNil)
260 //
261 // err = tree.PrintGraphviz(nil)
262 // c.Assert(err, qt.IsNil)
263}
264
265func TestGet(t *testing.T) {
266 c := qt.New(t)
267 database, err := db.NewBadgerDB(c.TempDir())
268 c.Assert(err, qt.IsNil)
269 tree, err := NewTree(database, 100, HashFunctionPoseidon)
270 c.Assert(err, qt.IsNil)
271 defer tree.db.Close() //nolint:errcheck
272
273 bLen := tree.HashFunction().Len()
274 for i := 0; i < 10; i++ {
275 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
276 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
277 if err := tree.Add(k, v); err != nil {
278 t.Fatal(err)
279 }
280 }
281
282 k := BigIntToBytes(bLen, big.NewInt(int64(7)))
283 gettedKey, gettedValue, err := tree.Get(k)
284 c.Assert(err, qt.IsNil)
285 c.Check(gettedKey, qt.DeepEquals, k)
286 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(int64(7*2))))
287}
288
289func TestGenProofAndVerify(t *testing.T) {
290 c := qt.New(t)
291 database, err := db.NewBadgerDB(c.TempDir())
292 c.Assert(err, qt.IsNil)
293 tree, err := NewTree(database, 100, HashFunctionPoseidon)
294 c.Assert(err, qt.IsNil)
295 defer tree.db.Close() //nolint:errcheck
296
297 bLen := tree.HashFunction().Len()
298 for i := 0; i < 10; i++ {
299 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
300 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
301 if err := tree.Add(k, v); err != nil {
302 t.Fatal(err)
303 }
304 }
305
306 k := BigIntToBytes(bLen, big.NewInt(int64(7)))
307 v := BigIntToBytes(bLen, big.NewInt(int64(14)))
308 kAux, proofV, siblings, existence, err := tree.GenProof(k)
309 c.Assert(err, qt.IsNil)
310 c.Assert(proofV, qt.DeepEquals, v)
311 c.Assert(k, qt.DeepEquals, kAux)
312 c.Assert(existence, qt.IsTrue)
313
314 verif, err := CheckProof(tree.hashFunction, k, v, tree.Root(), siblings)
315 c.Assert(err, qt.IsNil)
316 c.Check(verif, qt.IsTrue)
317}
318
319func TestDumpAndImportDump(t *testing.T) {
320 c := qt.New(t)
321 database1, err := db.NewBadgerDB(c.TempDir())
322 c.Assert(err, qt.IsNil)
323 tree1, err := NewTree(database1, 100, HashFunctionPoseidon)
324 c.Assert(err, qt.IsNil)
325 defer tree1.db.Close() //nolint:errcheck
326
327 bLen := tree1.HashFunction().Len()
328 for i := 0; i < 16; i++ {
329 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
330 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
331 if err := tree1.Add(k, v); err != nil {
332 t.Fatal(err)
333 }
334 }
335
336 e, err := tree1.Dump(nil)
337 c.Assert(err, qt.IsNil)
338
339 database2, err := db.NewBadgerDB(c.TempDir())
340 c.Assert(err, qt.IsNil)
341 tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
342 c.Assert(err, qt.IsNil)
343 defer tree2.db.Close() //nolint:errcheck
344 err = tree2.ImportDump(e)
345 c.Assert(err, qt.IsNil)
346 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
347 c.Check(hex.EncodeToString(tree2.Root()), qt.Equals,
348 "0d93aaa3362b2f999f15e15728f123087c2eee716f01c01f56e23aae07f09f08")
349}
350
351func TestRWMutex(t *testing.T) {
352 c := qt.New(t)
353 database, err := db.NewBadgerDB(c.TempDir())
354 c.Assert(err, qt.IsNil)
355 tree, err := NewTree(database, 100, HashFunctionPoseidon)
356 c.Assert(err, qt.IsNil)
357 defer tree.db.Close() //nolint:errcheck
358
359 bLen := tree.HashFunction().Len()
360 var keys, values [][]byte
361 for i := 0; i < 1000; i++ {
362 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
363 v := BigIntToBytes(bLen, big.NewInt(0))
364 keys = append(keys, k)
365 values = append(values, v)
366 }
367 go func() {
368 _, err = tree.AddBatch(keys, values)
369 if err != nil {
370 panic(err)
371 }
372 }()
373
374 time.Sleep(500 * time.Millisecond)
375 k := BigIntToBytes(bLen, big.NewInt(int64(99999)))
376 v := BigIntToBytes(bLen, big.NewInt(int64(99999)))
377 if err := tree.Add(k, v); err != nil {
378 t.Fatal(err)
379 }
380}
381
382func TestSetGetNLeafs(t *testing.T) {
383 c := qt.New(t)
384 database, err := db.NewBadgerDB(c.TempDir())
385 c.Assert(err, qt.IsNil)
386 tree, err := NewTree(database, 100, HashFunctionPoseidon)
387 c.Assert(err, qt.IsNil)
388
389 // 0
390 tree.dbBatch = tree.db.NewBatch()
391
392 err = tree.setNLeafs(0)
393 c.Assert(err, qt.IsNil)
394
395 err = tree.dbBatch.Write()
396 c.Assert(err, qt.IsNil)
397
398 n, err := tree.GetNLeafs()
399 c.Assert(err, qt.IsNil)
400 c.Assert(n, qt.Equals, 0)
401
402 // 1024
403 tree.dbBatch = tree.db.NewBatch()
404
405 err = tree.setNLeafs(1024)
406 c.Assert(err, qt.IsNil)
407
408 err = tree.dbBatch.Write()
409 c.Assert(err, qt.IsNil)
410
411 n, err = tree.GetNLeafs()
412 c.Assert(err, qt.IsNil)
413 c.Assert(n, qt.Equals, 1024)
414
415 // 2**64 -1
416 tree.dbBatch = tree.db.NewBatch()
417
418 maxUint := ^uint(0)
419 maxInt := int(maxUint >> 1)
420
421 err = tree.setNLeafs(maxInt)
422 c.Assert(err, qt.IsNil)
423
424 err = tree.dbBatch.Write()
425 c.Assert(err, qt.IsNil)
426
427 n, err = tree.GetNLeafs()
428 c.Assert(err, qt.IsNil)
429 c.Assert(n, qt.Equals, maxInt)
430}
431
432func TestSnapshot(t *testing.T) {
433 c := qt.New(t)
434 database, err := db.NewBadgerDB(c.TempDir())
435 c.Assert(err, qt.IsNil)
436 tree, err := NewTree(database, 100, HashFunctionPoseidon)
437 c.Assert(err, qt.IsNil)
438
439 // fill the tree
440 bLen := tree.HashFunction().Len()
441 var keys, values [][]byte
442 for i := 0; i < 1000; i++ {
443 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
444 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
445 keys = append(keys, k)
446 values = append(values, v)
447 }
448 indexes, err := tree.AddBatch(keys, values)
449 c.Assert(err, qt.IsNil)
450 c.Check(len(indexes), qt.Equals, 0)
451
452 rootBI := BytesToBigInt(tree.Root())
453 c.Check(rootBI.String(), qt.Equals,
454 "13742386369878513332697380582061714160370929283209286127733983161245560237407")
455
456 tree2, err := tree.Snapshot(nil)
457 c.Assert(err, qt.IsNil)
458 rootBI = BytesToBigInt(tree2.Root())
459 c.Check(rootBI.String(), qt.Equals,
460 "13742386369878513332697380582061714160370929283209286127733983161245560237407")
461
462 // add k-v to original tree
463 k := BigIntToBytes(bLen, big.NewInt(int64(1000)))
464 v := BigIntToBytes(bLen, big.NewInt(int64(1000)))
465 err = tree.Add(k, v)
466 c.Assert(err, qt.IsNil)
467
468 // expect original tree to have new root
469 rootBI = BytesToBigInt(tree.Root())
470 c.Check(rootBI.String(), qt.Equals,
471 "10747149055773881257049574592162159501044114324358186833013814735296193179713")
472
473 // expect snapshot tree to have the old root
474 rootBI = BytesToBigInt(tree2.Root())
475 c.Check(rootBI.String(), qt.Equals,
476 "13742386369878513332697380582061714160370929283209286127733983161245560237407")
477
478 err = tree2.Add(k, v)
479 c.Assert(err, qt.IsNil)
480 // after adding also the k-v into the snapshot tree, expect original
481 // tree to have new root
482 rootBI = BytesToBigInt(tree.Root())
483 c.Check(rootBI.String(), qt.Equals,
484 "10747149055773881257049574592162159501044114324358186833013814735296193179713")
485}
486
487func BenchmarkAdd(b *testing.B) {
488 bLen := 32 // for both Poseidon & Sha256
489 // prepare inputs
490 var ks, vs [][]byte
491 for i := 0; i < 1000; i++ {
492 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
493 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
494 ks = append(ks, k)
495 vs = append(vs, v)
496 }
497
498 b.Run("Poseidon", func(b *testing.B) {
499 benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
500 })
501 b.Run("Sha256", func(b *testing.B) {
502 benchmarkAdd(b, HashFunctionSha256, ks, vs)
503 })
504}
505
506func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
507 c := qt.New(b)
508 database, err := db.NewBadgerDB(c.TempDir())
509 c.Assert(err, qt.IsNil)
510 tree, err := NewTree(database, 140, hashFunc)
511 c.Assert(err, qt.IsNil)
512 defer tree.db.Close() //nolint:errcheck
513
514 for i := 0; i < len(ks); i++ {
515 if err := tree.Add(ks[i], vs[i]); err != nil {
516 b.Fatal(err)
517 }
518 }
519}