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