arbo/tree_test.go
tree_test.goRaw
1package arbo
2
3import (
4 "encoding/hex"
5 "math"
6 "math/big"
7 "runtime"
8 "testing"
9 "time"
10
11 qt "github.com/frankban/quicktest"
12 "go.vocdoni.io/dvote/db"
13 "go.vocdoni.io/dvote/db/badgerdb"
14)
15
16func checkRootBIString(c *qt.C, tree *Tree, expected string) {
17 root, err := tree.Root()
18 c.Assert(err, qt.IsNil)
19 rootBI := BytesToBigInt(root)
20 c.Check(rootBI.String(), qt.Equals, expected)
21}
22
23func TestDBTx(t *testing.T) {
24 c := qt.New(t)
25
26 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
27 c.Assert(err, qt.IsNil)
28 wTx := database.WriteTx()
29
30 _, err = wTx.Get([]byte("a"))
31 c.Assert(err, qt.Equals, db.ErrKeyNotFound)
32
33 err = wTx.Set([]byte("a"), []byte("b"))
34 c.Assert(err, qt.IsNil)
35
36 v, err := wTx.Get([]byte("a"))
37 c.Assert(err, qt.IsNil)
38 c.Assert(v, qt.DeepEquals, []byte("b"))
39}
40
41func TestAddTestVectors(t *testing.T) {
42 c := qt.New(t)
43
44 // Poseidon test vectors generated using https://github.com/iden3/circomlib smt.js
45 testVectorsPoseidon := []string{
46 "0000000000000000000000000000000000000000000000000000000000000000",
47 "13578938674299138072471463694055224830892726234048532520316387704878000008795",
48 "5412393676474193513566895793055462193090331607895808993925969873307089394741",
49 "14204494359367183802864593755198662203838502594566452929175967972147978322084",
50 }
51 testAdd(c, HashFunctionPoseidon, testVectorsPoseidon)
52
53 testVectorsSha256 := []string{
54 "0000000000000000000000000000000000000000000000000000000000000000",
55 "46910109172468462938850740851377282682950237270676610513794735904325820156367",
56 "59481735341404520835410489183267411392292882901306595567679529387376287440550",
57 "20573794434149960984975763118181266662429997821552560184909083010514790081771",
58 }
59 testAdd(c, HashFunctionSha256, testVectorsSha256)
60}
61
62func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
63 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
64 c.Assert(err, qt.IsNil)
65 tree, err := NewTree(database, 256, hashFunc)
66 c.Assert(err, qt.IsNil)
67 defer tree.db.Close() //nolint:errcheck
68
69 root, err := tree.Root()
70 c.Assert(err, qt.IsNil)
71 c.Check(hex.EncodeToString(root), qt.Equals, testVectors[0])
72
73 bLen := 32
74 err = tree.Add(
75 BigIntToBytes(bLen, big.NewInt(1)),
76 BigIntToBytes(bLen, big.NewInt(2)))
77 c.Assert(err, qt.IsNil)
78 checkRootBIString(c, tree, testVectors[1])
79
80 err = tree.Add(
81 BigIntToBytes(bLen, big.NewInt(33)),
82 BigIntToBytes(bLen, big.NewInt(44)))
83 c.Assert(err, qt.IsNil)
84 checkRootBIString(c, tree, testVectors[2])
85
86 err = tree.Add(
87 BigIntToBytes(bLen, big.NewInt(1234)),
88 BigIntToBytes(bLen, big.NewInt(9876)))
89 c.Assert(err, qt.IsNil)
90 checkRootBIString(c, tree, testVectors[3])
91}
92
93func TestAddBatch(t *testing.T) {
94 c := qt.New(t)
95 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
96 c.Assert(err, qt.IsNil)
97 tree, err := NewTree(database, 256, HashFunctionPoseidon)
98 c.Assert(err, qt.IsNil)
99 defer tree.db.Close() //nolint:errcheck
100
101 bLen := 32
102 for i := 0; i < 1000; i++ {
103 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
104 v := BigIntToBytes(bLen, big.NewInt(0))
105 if err := tree.Add(k, v); err != nil {
106 t.Fatal(err)
107 }
108 }
109
110 checkRootBIString(c, tree,
111 "296519252211642170490407814696803112091039265640052570497930797516015811235")
112
113 database, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()})
114 c.Assert(err, qt.IsNil)
115 tree2, err := NewTree(database, 256, HashFunctionPoseidon)
116 c.Assert(err, qt.IsNil)
117 defer tree2.db.Close() //nolint:errcheck
118
119 var keys, values [][]byte
120 for i := 0; i < 1000; i++ {
121 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
122 v := BigIntToBytes(bLen, big.NewInt(0))
123 keys = append(keys, k)
124 values = append(values, v)
125 }
126 indexes, err := tree2.AddBatch(keys, values)
127 c.Assert(err, qt.IsNil)
128 c.Check(len(indexes), qt.Equals, 0)
129
130 checkRootBIString(c, tree2,
131 "296519252211642170490407814696803112091039265640052570497930797516015811235")
132}
133
134func TestAddDifferentOrder(t *testing.T) {
135 c := qt.New(t)
136 database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
137 c.Assert(err, qt.IsNil)
138 tree1, err := NewTree(database1, 256, HashFunctionPoseidon)
139 c.Assert(err, qt.IsNil)
140 defer tree1.db.Close() //nolint:errcheck
141
142 bLen := 32
143 for i := 0; i < 16; i++ {
144 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
145 v := BigIntToBytes(bLen, big.NewInt(0))
146 if err := tree1.Add(k, v); err != nil {
147 t.Fatal(err)
148 }
149 }
150
151 database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
152 c.Assert(err, qt.IsNil)
153 tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
154 c.Assert(err, qt.IsNil)
155 defer tree2.db.Close() //nolint:errcheck
156
157 for i := 16 - 1; i >= 0; i-- {
158 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
159 v := BigIntToBytes(bLen, big.NewInt(0))
160 if err := tree2.Add(k, v); err != nil {
161 t.Fatal(err)
162 }
163 }
164
165 root1, err := tree1.Root()
166 c.Assert(err, qt.IsNil)
167 root2, err := tree2.Root()
168 c.Assert(err, qt.IsNil)
169 c.Check(hex.EncodeToString(root2), qt.Equals, hex.EncodeToString(root1))
170 c.Check(hex.EncodeToString(root1), qt.Equals,
171 "3b89100bec24da9275c87bc188740389e1d5accfc7d88ba5688d7fa96a00d82f")
172}
173
174func TestAddRepeatedIndex(t *testing.T) {
175 c := qt.New(t)
176 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
177 c.Assert(err, qt.IsNil)
178 tree, err := NewTree(database, 256, HashFunctionPoseidon)
179 c.Assert(err, qt.IsNil)
180 defer tree.db.Close() //nolint:errcheck
181
182 bLen := 32
183 k := BigIntToBytes(bLen, big.NewInt(int64(3)))
184 v := BigIntToBytes(bLen, big.NewInt(int64(12)))
185
186 err = tree.Add(k, v)
187 c.Assert(err, qt.IsNil)
188 err = tree.Add(k, v) // repeating same key-value
189 c.Check(err, qt.Equals, ErrKeyAlreadyExists)
190}
191
192func TestUpdate(t *testing.T) {
193 c := qt.New(t)
194 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
195 c.Assert(err, qt.IsNil)
196 tree, err := NewTree(database, 256, HashFunctionPoseidon)
197 c.Assert(err, qt.IsNil)
198 defer tree.db.Close() //nolint:errcheck
199
200 bLen := 32
201 k := BigIntToBytes(bLen, big.NewInt(int64(20)))
202 v := BigIntToBytes(bLen, big.NewInt(int64(12)))
203 if err := tree.Add(k, v); err != nil {
204 t.Fatal(err)
205 }
206
207 v = BigIntToBytes(bLen, big.NewInt(int64(11)))
208 err = tree.Update(k, v)
209 c.Assert(err, qt.IsNil)
210
211 gettedKey, gettedValue, err := tree.Get(k)
212 c.Assert(err, qt.IsNil)
213 c.Check(gettedKey, qt.DeepEquals, k)
214 c.Check(gettedValue, qt.DeepEquals, v)
215
216 // add more leafs to the tree to do another test
217 for i := 0; i < 16; i++ {
218 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
219 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
220 if err := tree.Add(k, v); err != nil {
221 t.Fatal(err)
222 }
223 }
224
225 k = BigIntToBytes(bLen, big.NewInt(int64(3)))
226 v = BigIntToBytes(bLen, big.NewInt(int64(11)))
227 // check that before the Update, 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.Not(qt.DeepEquals), v)
232 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(6)))
233
234 err = tree.Update(k, v)
235 c.Assert(err, qt.IsNil)
236
237 // check that after Update, the value for 3 is ==11
238 gettedKey, gettedValue, err = tree.Get(k)
239 c.Assert(err, qt.IsNil)
240 c.Check(gettedKey, qt.DeepEquals, k)
241 c.Check(gettedValue, qt.DeepEquals, v)
242 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(11)))
243}
244
245func TestAux(t *testing.T) { // TODO split in proper tests
246 c := qt.New(t)
247 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
248 c.Assert(err, qt.IsNil)
249 tree, err := NewTree(database, 256, HashFunctionPoseidon)
250 c.Assert(err, qt.IsNil)
251 defer tree.db.Close() //nolint:errcheck
252
253 bLen := 32
254 k := BigIntToBytes(bLen, big.NewInt(int64(1)))
255 v := BigIntToBytes(bLen, big.NewInt(int64(0)))
256 err = tree.Add(k, v)
257 c.Assert(err, qt.IsNil)
258 k = BigIntToBytes(bLen, big.NewInt(int64(256)))
259 err = tree.Add(k, v)
260 c.Assert(err, qt.IsNil)
261
262 k = BigIntToBytes(bLen, big.NewInt(int64(257)))
263 err = tree.Add(k, v)
264 c.Assert(err, qt.IsNil)
265
266 k = BigIntToBytes(bLen, big.NewInt(int64(515)))
267 err = tree.Add(k, v)
268 c.Assert(err, qt.IsNil)
269 k = BigIntToBytes(bLen, big.NewInt(int64(770)))
270 err = tree.Add(k, v)
271 c.Assert(err, qt.IsNil)
272
273 k = BigIntToBytes(bLen, big.NewInt(int64(388)))
274 err = tree.Add(k, v)
275 c.Assert(err, qt.IsNil)
276 k = BigIntToBytes(bLen, big.NewInt(int64(900)))
277 err = tree.Add(k, v)
278 c.Assert(err, qt.IsNil)
279 //
280 // err = tree.PrintGraphviz(nil)
281 // c.Assert(err, qt.IsNil)
282}
283
284func TestGet(t *testing.T) {
285 c := qt.New(t)
286 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
287 c.Assert(err, qt.IsNil)
288 tree, err := NewTree(database, 256, HashFunctionPoseidon)
289 c.Assert(err, qt.IsNil)
290 defer tree.db.Close() //nolint:errcheck
291
292 bLen := 32
293 for i := 0; i < 10; i++ {
294 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
295 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
296 if err := tree.Add(k, v); err != nil {
297 t.Fatal(err)
298 }
299 }
300
301 k := BigIntToBytes(bLen, big.NewInt(int64(7)))
302 gettedKey, gettedValue, err := tree.Get(k)
303 c.Assert(err, qt.IsNil)
304 c.Check(gettedKey, qt.DeepEquals, k)
305 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(int64(7*2))))
306}
307
308func TestGenProofAndVerify(t *testing.T) {
309 c := qt.New(t)
310 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
311 c.Assert(err, qt.IsNil)
312 tree, err := NewTree(database, 256, HashFunctionPoseidon)
313 c.Assert(err, qt.IsNil)
314 defer tree.db.Close() //nolint:errcheck
315
316 bLen := 32
317 for i := 0; i < 10; i++ {
318 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
319 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
320 if err := tree.Add(k, v); err != nil {
321 t.Fatal(err)
322 }
323 }
324
325 k := BigIntToBytes(bLen, big.NewInt(int64(7)))
326 v := BigIntToBytes(bLen, big.NewInt(int64(14)))
327 kAux, proofV, siblings, existence, err := tree.GenProof(k)
328 c.Assert(err, qt.IsNil)
329 c.Assert(proofV, qt.DeepEquals, v)
330 c.Assert(k, qt.DeepEquals, kAux)
331 c.Assert(existence, qt.IsTrue)
332
333 root, err := tree.Root()
334 c.Assert(err, qt.IsNil)
335 verif, err := CheckProof(tree.hashFunction, k, v, root, siblings)
336 c.Assert(err, qt.IsNil)
337 c.Check(verif, qt.IsTrue)
338}
339
340func TestDumpAndImportDump(t *testing.T) {
341 c := qt.New(t)
342 database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
343 c.Assert(err, qt.IsNil)
344 tree1, err := NewTree(database1, 256, HashFunctionPoseidon)
345 c.Assert(err, qt.IsNil)
346 defer tree1.db.Close() //nolint:errcheck
347
348 bLen := 32
349 for i := 0; i < 16; i++ {
350 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
351 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
352 if err := tree1.Add(k, v); err != nil {
353 t.Fatal(err)
354 }
355 }
356
357 e, err := tree1.Dump(nil)
358 c.Assert(err, qt.IsNil)
359
360 database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
361 c.Assert(err, qt.IsNil)
362 tree2, err := NewTree(database2, 256, HashFunctionPoseidon)
363 c.Assert(err, qt.IsNil)
364 defer tree2.db.Close() //nolint:errcheck
365 err = tree2.ImportDump(e)
366 c.Assert(err, qt.IsNil)
367
368 root1, err := tree1.Root()
369 c.Assert(err, qt.IsNil)
370 root2, err := tree2.Root()
371 c.Assert(err, qt.IsNil)
372 c.Check(root2, qt.DeepEquals, root1)
373 c.Check(hex.EncodeToString(root2), qt.Equals,
374 "0d93aaa3362b2f999f15e15728f123087c2eee716f01c01f56e23aae07f09f08")
375}
376
377func TestRWMutex(t *testing.T) {
378 c := qt.New(t)
379 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
380 c.Assert(err, qt.IsNil)
381 tree, err := NewTree(database, 256, HashFunctionPoseidon)
382 c.Assert(err, qt.IsNil)
383 defer tree.db.Close() //nolint:errcheck
384
385 bLen := 32
386 var keys, values [][]byte
387 for i := 0; i < 1000; i++ {
388 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
389 v := BigIntToBytes(bLen, big.NewInt(0))
390 keys = append(keys, k)
391 values = append(values, v)
392 }
393 go func() {
394 _, err = tree.AddBatch(keys, values)
395 if err != nil {
396 panic(err)
397 }
398 }()
399
400 time.Sleep(500 * time.Millisecond)
401 k := BigIntToBytes(bLen, big.NewInt(int64(99999)))
402 v := BigIntToBytes(bLen, big.NewInt(int64(99999)))
403 if err := tree.Add(k, v); err != nil {
404 t.Fatal(err)
405 }
406}
407
408// TODO UPDATE
409// func TestSetGetNLeafs(t *testing.T) {
410// c := qt.New(t)
411// database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
412// c.Assert(err, qt.IsNil)
413// tree, err := NewTree(database, 100, HashFunctionPoseidon)
414// c.Assert(err, qt.IsNil)
415//
416// // 0
417// tree.dbBatch = tree.db.NewBatch()
418//
419// err = tree.setNLeafs(0)
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, 0)
428//
429// // 1024
430// tree.dbBatch = tree.db.NewBatch()
431//
432// err = tree.setNLeafs(1024)
433// c.Assert(err, qt.IsNil)
434//
435// err = tree.dbBatch.Write()
436// c.Assert(err, qt.IsNil)
437//
438// n, err = tree.GetNLeafs()
439// c.Assert(err, qt.IsNil)
440// c.Assert(n, qt.Equals, 1024)
441//
442// // 2**64 -1
443// tree.dbBatch = tree.db.NewBatch()
444//
445// maxUint := ^uint(0)
446// maxInt := int(maxUint >> 1)
447//
448// err = tree.setNLeafs(maxInt)
449// c.Assert(err, qt.IsNil)
450//
451// err = tree.dbBatch.Write()
452// c.Assert(err, qt.IsNil)
453//
454// n, err = tree.GetNLeafs()
455// c.Assert(err, qt.IsNil)
456// c.Assert(n, qt.Equals, maxInt)
457// }
458
459func TestAddBatchFullyUsed(t *testing.T) {
460 c := qt.New(t)
461
462 database1, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
463 c.Assert(err, qt.IsNil)
464 tree1, err := NewTree(database1, 4, HashFunctionPoseidon)
465 c.Assert(err, qt.IsNil)
466
467 database2, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
468 c.Assert(err, qt.IsNil)
469 tree2, err := NewTree(database2, 4, HashFunctionPoseidon)
470 c.Assert(err, qt.IsNil)
471
472 var keys, values [][]byte
473 for i := 0; i < 16; i++ {
474 k := BigIntToBytes(1, big.NewInt(int64(i)))
475 v := k
476
477 keys = append(keys, k)
478 values = append(values, v)
479
480 // add one by one expecting no error
481 err := tree1.Add(k, v)
482 c.Assert(err, qt.IsNil)
483 }
484
485 invalids, err := tree2.AddBatch(keys, values)
486 c.Assert(err, qt.IsNil)
487 c.Assert(0, qt.Equals, len(invalids))
488
489 root1, err := tree1.Root()
490 c.Assert(err, qt.IsNil)
491 root2, err := tree2.Root()
492 c.Assert(err, qt.IsNil)
493 c.Assert(root1, qt.DeepEquals, root2)
494
495 // get all key-values and check that are equal between both trees
496 for i := 0; i < 16; i++ {
497 auxK1, auxV1, err := tree1.Get(BigIntToBytes(1, big.NewInt(int64(i))))
498 c.Assert(err, qt.IsNil)
499
500 auxK2, auxV2, err := tree2.Get(BigIntToBytes(1, big.NewInt(int64(i))))
501 c.Assert(err, qt.IsNil)
502
503 c.Assert(auxK1, qt.DeepEquals, auxK2)
504 c.Assert(auxV1, qt.DeepEquals, auxV2)
505 }
506
507 // try adding one more key to both trees (through Add & AddBatch) and
508 // expect not being added due the tree is already full
509 k := BigIntToBytes(1, big.NewInt(int64(16)))
510 v := k
511 err = tree1.Add(k, v)
512 c.Assert(err, qt.Equals, ErrMaxVirtualLevel)
513
514 invalids, err = tree2.AddBatch([][]byte{k}, [][]byte{v})
515 c.Assert(err, qt.IsNil)
516 c.Assert(1, qt.Equals, len(invalids))
517}
518
519func TestSetRoot(t *testing.T) {
520 c := qt.New(t)
521 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
522 c.Assert(err, qt.IsNil)
523 tree, err := NewTree(database, 256, HashFunctionPoseidon)
524 c.Assert(err, qt.IsNil)
525
526 expectedRoot := "13742386369878513332697380582061714160370929283209286127733983161245560237407"
527
528 // fill the tree
529 bLen := 32
530 var keys, values [][]byte
531 for i := 0; i < 1000; i++ {
532 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
533 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
534 keys = append(keys, k)
535 values = append(values, v)
536 }
537 indexes, err := tree.AddBatch(keys, values)
538 c.Assert(err, qt.IsNil)
539 c.Check(len(indexes), qt.Equals, 0)
540 checkRootBIString(c, tree,
541 expectedRoot)
542
543 // add one more k-v
544 k := BigIntToBytes(bLen, big.NewInt(1000))
545 v := BigIntToBytes(bLen, big.NewInt(1000))
546 err = tree.Add(k, v)
547 c.Assert(err, qt.IsNil)
548 checkRootBIString(c, tree,
549 "10747149055773881257049574592162159501044114324358186833013814735296193179713")
550
551 // do a SetRoot, and expect the same root than the original tree
552 pastRootBI, ok := new(big.Int).SetString(expectedRoot, 10)
553 c.Assert(ok, qt.IsTrue)
554 pastRoot := BigIntToBytes(32, pastRootBI)
555
556 err = tree.SetRoot(pastRoot)
557 c.Assert(err, qt.IsNil)
558 checkRootBIString(c, tree, expectedRoot)
559
560 // check that the tree can be updated
561 err = tree.Add([]byte("test"), []byte("test"))
562 c.Assert(err, qt.IsNil)
563 err = tree.Update([]byte("test"), []byte("test"))
564 c.Assert(err, qt.IsNil)
565
566 // check that the k-v '1000' does not exist in the new tree
567 _, _, err = tree.Get(k)
568 c.Assert(err, qt.Equals, ErrKeyNotFound)
569
570 // check that can be set an empty root
571 err = tree.SetRoot(tree.emptyHash)
572 c.Assert(err, qt.IsNil)
573}
574
575func TestSnapshot(t *testing.T) {
576 c := qt.New(t)
577 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
578 c.Assert(err, qt.IsNil)
579 tree, err := NewTree(database, 256, HashFunctionPoseidon)
580 c.Assert(err, qt.IsNil)
581
582 // fill the tree
583 bLen := 32
584 var keys, values [][]byte
585 for i := 0; i < 1000; i++ {
586 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
587 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
588 keys = append(keys, k)
589 values = append(values, v)
590 }
591 indexes, err := tree.AddBatch(keys, values)
592 c.Assert(err, qt.IsNil)
593 c.Check(len(indexes), qt.Equals, 0)
594 checkRootBIString(c, tree,
595 "13742386369878513332697380582061714160370929283209286127733983161245560237407")
596
597 // do a snapshot, and expect the same root than the original tree
598 snapshotTree, err := tree.Snapshot(nil)
599 c.Assert(err, qt.IsNil)
600 checkRootBIString(c, snapshotTree,
601 "13742386369878513332697380582061714160370929283209286127733983161245560237407")
602
603 // check that the snapshotTree can not be updated
604 _, err = snapshotTree.AddBatch(keys, values)
605 c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
606 err = snapshotTree.Add([]byte("test"), []byte("test"))
607 c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
608 err = snapshotTree.Update([]byte("test"), []byte("test"))
609 c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
610 err = snapshotTree.ImportDump(nil)
611 c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
612
613 // update the original tree by adding a new key-value, and check that
614 // snapshotTree still has the old root, but the original tree has a new
615 // root
616 err = tree.Add([]byte("test"), []byte("test"))
617 c.Assert(err, qt.IsNil)
618 checkRootBIString(c, snapshotTree,
619 "13742386369878513332697380582061714160370929283209286127733983161245560237407")
620 checkRootBIString(c, tree,
621 "1025190963769001718196479367844646783678188389989148142691917685159698888868")
622}
623
624func TestGetFromSnapshotExpectArboErrKeyNotFound(t *testing.T) {
625 c := qt.New(t)
626
627 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
628 c.Assert(err, qt.IsNil)
629 tree, err := NewTree(database, 256, HashFunctionPoseidon)
630 c.Assert(err, qt.IsNil)
631 defer tree.db.Close() //nolint:errcheck
632
633 bLen := 32
634 k := BigIntToBytes(bLen, big.NewInt(int64(3)))
635
636 root, err := tree.Root()
637 c.Assert(err, qt.IsNil)
638 tree, err = tree.Snapshot(root)
639 c.Assert(err, qt.IsNil)
640
641 _, _, err = tree.Get(k)
642 c.Assert(err, qt.Equals, ErrKeyNotFound) // and not equal to db.ErrKeyNotFound
643}
644
645func TestKeyLen(t *testing.T) {
646 c := qt.New(t)
647 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
648 c.Assert(err, qt.IsNil)
649 // maxLevels is 100, keyPath length = ceil(maxLevels/8) = 13
650 maxLevels := 100
651 tree, err := NewTree(database, maxLevels, HashFunctionBlake2b)
652 c.Assert(err, qt.IsNil)
653
654 // expect no errors when adding a key of only 4 bytes (when the
655 // required length of keyPath for 100 levels would be 13 bytes)
656 bLen := 4
657 k := BigIntToBytes(bLen, big.NewInt(1))
658 v := BigIntToBytes(bLen, big.NewInt(1))
659
660 err = tree.Add(k, v)
661 c.Assert(err, qt.IsNil)
662
663 err = tree.Update(k, v)
664 c.Assert(err, qt.IsNil)
665
666 _, _, _, _, err = tree.GenProof(k)
667 c.Assert(err, qt.IsNil)
668
669 _, _, err = tree.Get(k)
670 c.Assert(err, qt.IsNil)
671
672 k = BigIntToBytes(bLen, big.NewInt(2))
673 v = BigIntToBytes(bLen, big.NewInt(2))
674 invalids, err := tree.AddBatch([][]byte{k}, [][]byte{v})
675 c.Assert(err, qt.IsNil)
676 c.Assert(len(invalids), qt.Equals, 0)
677
678 // expect errors when adding a key bigger than maximum capacity of the
679 // tree: ceil(maxLevels/8)
680 maxLevels = 32
681 database, err = badgerdb.New(badgerdb.Options{Path: c.TempDir()})
682 c.Assert(err, qt.IsNil)
683 tree, err = NewTree(database, maxLevels, HashFunctionBlake2b)
684 c.Assert(err, qt.IsNil)
685
686 maxKeyLen := int(math.Ceil(float64(maxLevels) / float64(8))) //nolint:gomnd
687 k = BigIntToBytes(maxKeyLen+1, big.NewInt(1))
688 v = BigIntToBytes(maxKeyLen+1, big.NewInt(1))
689
690 expectedErrMsg := "len(k) can not be bigger than ceil(maxLevels/8)," +
691 " where len(k): 5, maxLevels: 32, max key len=ceil(maxLevels/8): 4." +
692 " Might need a bigger tree depth (maxLevels>=40) in order to input" +
693 " keys of length 5"
694
695 err = tree.Add(k, v)
696 c.Assert(err.Error(), qt.Equals, expectedErrMsg)
697
698 err = tree.Update(k, v)
699 c.Assert(err.Error(), qt.Equals, expectedErrMsg)
700
701 _, _, _, _, err = tree.GenProof(k)
702 c.Assert(err.Error(), qt.Equals, expectedErrMsg)
703
704 _, _, err = tree.Get(k)
705 c.Assert(err.Error(), qt.Equals, expectedErrMsg)
706
707 // check AddBatch with few key-values
708 invalids, err = tree.AddBatch([][]byte{k}, [][]byte{v})
709 c.Assert(err, qt.IsNil)
710 c.Assert(len(invalids), qt.Equals, 1)
711
712 // check AddBatch with many key-values
713 nCPU := flp2(runtime.NumCPU())
714 nKVs := nCPU + 1
715 var ks, vs [][]byte
716 for i := 0; i < nKVs; i++ {
717 ks = append(ks, BigIntToBytes(maxKeyLen+1, big.NewInt(1)))
718 vs = append(vs, BigIntToBytes(maxKeyLen+1, big.NewInt(1)))
719 }
720 invalids, err = tree.AddBatch(ks, vs)
721 c.Assert(err, qt.IsNil)
722 c.Assert(len(invalids), qt.Equals, nKVs)
723
724 // check that with maxKeyLen it can be added
725 k = BigIntToBytes(maxKeyLen, big.NewInt(1))
726 err = tree.Add(k, v)
727 c.Assert(err, qt.IsNil)
728
729 // check CheckProof check with key longer
730 kAux, vAux, packedSiblings, existence, err := tree.GenProof(k)
731 c.Assert(err, qt.IsNil)
732 c.Assert(existence, qt.IsTrue)
733
734 root, err := tree.Root()
735 c.Assert(err, qt.IsNil)
736 verif, err := CheckProof(tree.HashFunction(), kAux, vAux, root, packedSiblings)
737 c.Assert(err, qt.IsNil)
738 c.Assert(verif, qt.IsTrue)
739
740 // use a similar key but with one zero, expect that CheckProof fails on
741 // the verification
742 kAux = append(kAux, 0)
743 verif, err = CheckProof(tree.HashFunction(), kAux, vAux, root, packedSiblings)
744 c.Assert(err, qt.IsNil)
745 c.Assert(verif, qt.IsFalse)
746}
747
748func BenchmarkAdd(b *testing.B) {
749 bLen := 32 // for both Poseidon & Sha256
750 // prepare inputs
751 var ks, vs [][]byte
752 for i := 0; i < 1000; i++ {
753 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
754 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
755 ks = append(ks, k)
756 vs = append(vs, v)
757 }
758
759 b.Run("Poseidon", func(b *testing.B) {
760 benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
761 })
762 b.Run("Sha256", func(b *testing.B) {
763 benchmarkAdd(b, HashFunctionSha256, ks, vs)
764 })
765}
766
767func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
768 c := qt.New(b)
769 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
770 c.Assert(err, qt.IsNil)
771 tree, err := NewTree(database, 140, hashFunc)
772 c.Assert(err, qt.IsNil)
773 defer tree.db.Close() //nolint:errcheck
774
775 for i := 0; i < len(ks); i++ {
776 if err := tree.Add(ks[i], vs[i]); err != nil {
777 b.Fatal(err)
778 }
779 }
780}