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 proofV, siblings, err := tree.GenProof(k)
309 c.Assert(err, qt.IsNil)
310 c.Assert(proofV, qt.DeepEquals, v)
311
312 verif, err := CheckProof(tree.hashFunction, k, v, tree.Root(), siblings)
313 c.Assert(err, qt.IsNil)
314 c.Check(verif, qt.IsTrue)
315}
316
317func TestDumpAndImportDump(t *testing.T) {
318 c := qt.New(t)
319 database1, err := db.NewBadgerDB(c.TempDir())
320 c.Assert(err, qt.IsNil)
321 tree1, err := NewTree(database1, 100, HashFunctionPoseidon)
322 c.Assert(err, qt.IsNil)
323 defer tree1.db.Close() //nolint:errcheck
324
325 bLen := tree1.HashFunction().Len()
326 for i := 0; i < 16; i++ {
327 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
328 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
329 if err := tree1.Add(k, v); err != nil {
330 t.Fatal(err)
331 }
332 }
333
334 e, err := tree1.Dump(nil)
335 c.Assert(err, qt.IsNil)
336
337 database2, err := db.NewBadgerDB(c.TempDir())
338 c.Assert(err, qt.IsNil)
339 tree2, err := NewTree(database2, 100, HashFunctionPoseidon)
340 c.Assert(err, qt.IsNil)
341 defer tree2.db.Close() //nolint:errcheck
342 err = tree2.ImportDump(e)
343 c.Assert(err, qt.IsNil)
344 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
345 c.Check(hex.EncodeToString(tree2.Root()), qt.Equals,
346 "0d93aaa3362b2f999f15e15728f123087c2eee716f01c01f56e23aae07f09f08")
347}
348
349func TestRWMutex(t *testing.T) {
350 c := qt.New(t)
351 database, err := db.NewBadgerDB(c.TempDir())
352 c.Assert(err, qt.IsNil)
353 tree, err := NewTree(database, 100, HashFunctionPoseidon)
354 c.Assert(err, qt.IsNil)
355 defer tree.db.Close() //nolint:errcheck
356
357 bLen := tree.HashFunction().Len()
358 var keys, values [][]byte
359 for i := 0; i < 1000; i++ {
360 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
361 v := BigIntToBytes(bLen, big.NewInt(0))
362 keys = append(keys, k)
363 values = append(values, v)
364 }
365 go func() {
366 _, err = tree.AddBatch(keys, values)
367 if err != nil {
368 panic(err)
369 }
370 }()
371
372 time.Sleep(500 * time.Millisecond)
373 k := BigIntToBytes(bLen, big.NewInt(int64(99999)))
374 v := BigIntToBytes(bLen, big.NewInt(int64(99999)))
375 if err := tree.Add(k, v); err != nil {
376 t.Fatal(err)
377 }
378}
379
380func TestSetGetNLeafs(t *testing.T) {
381 c := qt.New(t)
382 database, err := db.NewBadgerDB(c.TempDir())
383 c.Assert(err, qt.IsNil)
384 tree, err := NewTree(database, 100, HashFunctionPoseidon)
385 c.Assert(err, qt.IsNil)
386
387 // 0
388 tree.dbBatch = tree.db.NewBatch()
389
390 err = tree.setNLeafs(0)
391 c.Assert(err, qt.IsNil)
392
393 err = tree.dbBatch.Write()
394 c.Assert(err, qt.IsNil)
395
396 n, err := tree.GetNLeafs()
397 c.Assert(err, qt.IsNil)
398 c.Assert(n, qt.Equals, 0)
399
400 // 1024
401 tree.dbBatch = tree.db.NewBatch()
402
403 err = tree.setNLeafs(1024)
404 c.Assert(err, qt.IsNil)
405
406 err = tree.dbBatch.Write()
407 c.Assert(err, qt.IsNil)
408
409 n, err = tree.GetNLeafs()
410 c.Assert(err, qt.IsNil)
411 c.Assert(n, qt.Equals, 1024)
412
413 // 2**64 -1
414 tree.dbBatch = tree.db.NewBatch()
415
416 maxUint := ^uint(0)
417 maxInt := int(maxUint >> 1)
418
419 err = tree.setNLeafs(maxInt)
420 c.Assert(err, qt.IsNil)
421
422 err = tree.dbBatch.Write()
423 c.Assert(err, qt.IsNil)
424
425 n, err = tree.GetNLeafs()
426 c.Assert(err, qt.IsNil)
427 c.Assert(n, qt.Equals, maxInt)
428}
429
430func BenchmarkAdd(b *testing.B) {
431 bLen := 32 // for both Poseidon & Sha256
432 // prepare inputs
433 var ks, vs [][]byte
434 for i := 0; i < 1000; i++ {
435 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
436 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
437 ks = append(ks, k)
438 vs = append(vs, v)
439 }
440
441 b.Run("Poseidon", func(b *testing.B) {
442 benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
443 })
444 b.Run("Sha256", func(b *testing.B) {
445 benchmarkAdd(b, HashFunctionSha256, ks, vs)
446 })
447}
448
449func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
450 c := qt.New(b)
451 database, err := db.NewBadgerDB(c.TempDir())
452 c.Assert(err, qt.IsNil)
453 tree, err := NewTree(database, 140, hashFunc)
454 c.Assert(err, qt.IsNil)
455 defer tree.db.Close() //nolint:errcheck
456
457 for i := 0; i < len(ks); i++ {
458 if err := tree.Add(ks[i], vs[i]); err != nil {
459 b.Fatal(err)
460 }
461 }
462}