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