arbo/tree_test.go
tree_test.goRaw
1package arbo
2
3import (
4 "bufio"
5 "encoding/hex"
6 "math"
7 "math/big"
8 "os"
9 "path/filepath"
10 "runtime"
11 "testing"
12 "time"
13
14 qt "github.com/frankban/quicktest"
15 "go.vocdoni.io/dvote/db"
16 "go.vocdoni.io/dvote/db/badgerdb"
17 "go.vocdoni.io/dvote/db/pebbledb"
18)
19
20func checkRootBIString(c *qt.C, tree *Tree, expected string) {
21 root, err := tree.Root()
22 c.Assert(err, qt.IsNil)
23 rootBI := BytesToBigInt(root)
24 c.Check(rootBI.String(), qt.Equals, expected)
25}
26
27func TestDBTx(t *testing.T) {
28 c := qt.New(t)
29
30 dbBadger, err := badgerdb.New(db.Options{Path: c.TempDir()})
31 c.Assert(err, qt.IsNil)
32 testDBTx(c, dbBadger)
33
34 dbPebble, err := pebbledb.New(db.Options{Path: c.TempDir()})
35 c.Assert(err, qt.IsNil)
36 testDBTx(c, dbPebble)
37}
38
39func testDBTx(c *qt.C, database db.Database) {
40 wTx := database.WriteTx()
41 defer wTx.Discard()
42
43 _, err := wTx.Get([]byte("a"))
44 c.Assert(err, qt.Equals, db.ErrKeyNotFound)
45
46 err = wTx.Set([]byte("a"), []byte("b"))
47 c.Assert(err, qt.IsNil)
48
49 v, err := wTx.Get([]byte("a"))
50 c.Assert(err, qt.IsNil)
51 c.Assert(v, qt.DeepEquals, []byte("b"))
52}
53
54func TestAddTestVectors(t *testing.T) {
55 c := qt.New(t)
56
57 // Poseidon test vectors generated using https://github.com/iden3/circomlib smt.js
58 testVectorsPoseidon := []string{
59 "0000000000000000000000000000000000000000000000000000000000000000",
60 "13578938674299138072471463694055224830892726234048532520316387704878000008795",
61 "5412393676474193513566895793055462193090331607895808993925969873307089394741",
62 "14204494359367183802864593755198662203838502594566452929175967972147978322084",
63 }
64 testAdd(c, HashFunctionPoseidon, testVectorsPoseidon)
65
66 testVectorsSha256 := []string{
67 "0000000000000000000000000000000000000000000000000000000000000000",
68 "46910109172468462938850740851377282682950237270676610513794735904325820156367",
69 "59481735341404520835410489183267411392292882901306595567679529387376287440550",
70 "20573794434149960984975763118181266662429997821552560184909083010514790081771",
71 }
72 testAdd(c, HashFunctionSha256, testVectorsSha256)
73}
74
75func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
76 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
77 c.Assert(err, qt.IsNil)
78 tree, err := NewTree(Config{Database: database, MaxLevels: 256,
79 HashFunction: hashFunc})
80 c.Assert(err, qt.IsNil)
81 defer tree.db.Close() //nolint:errcheck
82
83 root, err := tree.Root()
84 c.Assert(err, qt.IsNil)
85 c.Check(hex.EncodeToString(root), qt.Equals, testVectors[0])
86
87 bLen := 32
88 err = tree.Add(
89 BigIntToBytes(bLen, big.NewInt(1)),
90 BigIntToBytes(bLen, big.NewInt(2)))
91 c.Assert(err, qt.IsNil)
92 checkRootBIString(c, tree, testVectors[1])
93
94 err = tree.Add(
95 BigIntToBytes(bLen, big.NewInt(33)),
96 BigIntToBytes(bLen, big.NewInt(44)))
97 c.Assert(err, qt.IsNil)
98 checkRootBIString(c, tree, testVectors[2])
99
100 err = tree.Add(
101 BigIntToBytes(bLen, big.NewInt(1234)),
102 BigIntToBytes(bLen, big.NewInt(9876)))
103 c.Assert(err, qt.IsNil)
104 checkRootBIString(c, tree, testVectors[3])
105}
106
107func TestAddBatch(t *testing.T) {
108 c := qt.New(t)
109 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
110 c.Assert(err, qt.IsNil)
111 tree, err := NewTree(Config{Database: database, MaxLevels: 256,
112 HashFunction: HashFunctionPoseidon})
113 c.Assert(err, qt.IsNil)
114 defer tree.db.Close() //nolint:errcheck
115
116 bLen := 32
117 for i := 0; i < 1000; i++ {
118 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
119 v := BigIntToBytes(bLen, big.NewInt(0))
120 if err := tree.Add(k, v); err != nil {
121 t.Fatal(err)
122 }
123 }
124
125 checkRootBIString(c, tree,
126 "296519252211642170490407814696803112091039265640052570497930797516015811235")
127
128 database, err = badgerdb.New(db.Options{Path: c.TempDir()})
129 c.Assert(err, qt.IsNil)
130 tree2, err := NewTree(Config{Database: database, MaxLevels: 256,
131 HashFunction: HashFunctionPoseidon})
132 c.Assert(err, qt.IsNil)
133 defer tree2.db.Close() //nolint:errcheck
134
135 var keys, values [][]byte
136 for i := 0; i < 1000; i++ {
137 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
138 v := BigIntToBytes(bLen, big.NewInt(0))
139 keys = append(keys, k)
140 values = append(values, v)
141 }
142 indexes, err := tree2.AddBatch(keys, values)
143 c.Assert(err, qt.IsNil)
144 c.Check(len(indexes), qt.Equals, 0)
145
146 checkRootBIString(c, tree2,
147 "296519252211642170490407814696803112091039265640052570497930797516015811235")
148}
149
150func TestAddDifferentOrder(t *testing.T) {
151 c := qt.New(t)
152 database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
153 c.Assert(err, qt.IsNil)
154 tree1, err := NewTree(Config{Database: database1, MaxLevels: 256,
155 HashFunction: HashFunctionPoseidon})
156 c.Assert(err, qt.IsNil)
157 defer tree1.db.Close() //nolint:errcheck
158
159 bLen := 32
160 for i := 0; i < 16; i++ {
161 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
162 v := BigIntToBytes(bLen, big.NewInt(0))
163 if err := tree1.Add(k, v); err != nil {
164 t.Fatal(err)
165 }
166 }
167
168 database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
169 c.Assert(err, qt.IsNil)
170 tree2, err := NewTree(Config{Database: database2, MaxLevels: 256,
171 HashFunction: HashFunctionPoseidon})
172 c.Assert(err, qt.IsNil)
173 defer tree2.db.Close() //nolint:errcheck
174
175 for i := 16 - 1; i >= 0; i-- {
176 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
177 v := BigIntToBytes(bLen, big.NewInt(0))
178 if err := tree2.Add(k, v); err != nil {
179 t.Fatal(err)
180 }
181 }
182
183 root1, err := tree1.Root()
184 c.Assert(err, qt.IsNil)
185 root2, err := tree2.Root()
186 c.Assert(err, qt.IsNil)
187 c.Check(hex.EncodeToString(root2), qt.Equals, hex.EncodeToString(root1))
188 c.Check(hex.EncodeToString(root1), qt.Equals,
189 "3b89100bec24da9275c87bc188740389e1d5accfc7d88ba5688d7fa96a00d82f")
190}
191
192func TestAddRepeatedIndex(t *testing.T) {
193 c := qt.New(t)
194 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
195 c.Assert(err, qt.IsNil)
196 tree, err := NewTree(Config{Database: database, MaxLevels: 256,
197 HashFunction: HashFunctionPoseidon})
198 c.Assert(err, qt.IsNil)
199 defer tree.db.Close() //nolint:errcheck
200
201 bLen := 32
202 k := BigIntToBytes(bLen, big.NewInt(int64(3)))
203 v := BigIntToBytes(bLen, big.NewInt(int64(12)))
204
205 err = tree.Add(k, v)
206 c.Assert(err, qt.IsNil)
207 err = tree.Add(k, v) // repeating same key-value
208 c.Check(err, qt.Equals, ErrKeyAlreadyExists)
209}
210
211func TestUpdate(t *testing.T) {
212 c := qt.New(t)
213 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
214 c.Assert(err, qt.IsNil)
215 tree, err := NewTree(Config{Database: database, MaxLevels: 256,
216 HashFunction: HashFunctionPoseidon})
217 c.Assert(err, qt.IsNil)
218 defer tree.db.Close() //nolint:errcheck
219
220 bLen := 32
221 k := BigIntToBytes(bLen, big.NewInt(int64(20)))
222 v := BigIntToBytes(bLen, big.NewInt(int64(12)))
223 if err := tree.Add(k, v); err != nil {
224 t.Fatal(err)
225 }
226
227 v = BigIntToBytes(bLen, big.NewInt(int64(11)))
228 err = tree.Update(k, v)
229 c.Assert(err, qt.IsNil)
230
231 gettedKey, gettedValue, err := tree.Get(k)
232 c.Assert(err, qt.IsNil)
233 c.Check(gettedKey, qt.DeepEquals, k)
234 c.Check(gettedValue, qt.DeepEquals, v)
235
236 // add more leafs to the tree to do another test
237 for i := 0; i < 16; i++ {
238 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
239 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
240 if err := tree.Add(k, v); err != nil {
241 t.Fatal(err)
242 }
243 }
244
245 k = BigIntToBytes(bLen, big.NewInt(int64(3)))
246 v = BigIntToBytes(bLen, big.NewInt(int64(11)))
247 // check that before the Update, value for 3 is !=11
248 gettedKey, gettedValue, err = tree.Get(k)
249 c.Assert(err, qt.IsNil)
250 c.Check(gettedKey, qt.DeepEquals, k)
251 c.Check(gettedValue, qt.Not(qt.DeepEquals), v)
252 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(6)))
253
254 err = tree.Update(k, v)
255 c.Assert(err, qt.IsNil)
256
257 // check that after Update, the value for 3 is ==11
258 gettedKey, gettedValue, err = tree.Get(k)
259 c.Assert(err, qt.IsNil)
260 c.Check(gettedKey, qt.DeepEquals, k)
261 c.Check(gettedValue, qt.DeepEquals, v)
262 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(11)))
263}
264
265func TestAux(t *testing.T) { // TODO split in proper tests
266 c := qt.New(t)
267 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
268 c.Assert(err, qt.IsNil)
269 tree, err := NewTree(Config{Database: database, MaxLevels: 256,
270 HashFunction: HashFunctionPoseidon})
271 c.Assert(err, qt.IsNil)
272 defer tree.db.Close() //nolint:errcheck
273
274 bLen := 32
275 k := BigIntToBytes(bLen, big.NewInt(int64(1)))
276 v := BigIntToBytes(bLen, big.NewInt(int64(0)))
277 err = tree.Add(k, v)
278 c.Assert(err, qt.IsNil)
279 k = BigIntToBytes(bLen, big.NewInt(int64(256)))
280 err = tree.Add(k, v)
281 c.Assert(err, qt.IsNil)
282
283 k = BigIntToBytes(bLen, big.NewInt(int64(257)))
284 err = tree.Add(k, v)
285 c.Assert(err, qt.IsNil)
286
287 k = BigIntToBytes(bLen, big.NewInt(int64(515)))
288 err = tree.Add(k, v)
289 c.Assert(err, qt.IsNil)
290 k = BigIntToBytes(bLen, big.NewInt(int64(770)))
291 err = tree.Add(k, v)
292 c.Assert(err, qt.IsNil)
293
294 k = BigIntToBytes(bLen, big.NewInt(int64(388)))
295 err = tree.Add(k, v)
296 c.Assert(err, qt.IsNil)
297 k = BigIntToBytes(bLen, big.NewInt(int64(900)))
298 err = tree.Add(k, v)
299 c.Assert(err, qt.IsNil)
300 //
301 // err = tree.PrintGraphviz(nil)
302 // c.Assert(err, qt.IsNil)
303}
304
305func TestGet(t *testing.T) {
306 c := qt.New(t)
307 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
308 c.Assert(err, qt.IsNil)
309 tree, err := NewTree(Config{Database: database, MaxLevels: 256,
310 HashFunction: HashFunctionPoseidon})
311 c.Assert(err, qt.IsNil)
312 defer tree.db.Close() //nolint:errcheck
313
314 bLen := 32
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 gettedKey, gettedValue, err := tree.Get(k)
325 c.Assert(err, qt.IsNil)
326 c.Check(gettedKey, qt.DeepEquals, k)
327 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(int64(7*2))))
328}
329
330func TestBitmapBytes(t *testing.T) {
331 c := qt.New(t)
332
333 b := []byte{15}
334 bits := bytesToBitmap(b)
335 c.Assert(bits, qt.DeepEquals, []bool{true, true, true, true,
336 false, false, false, false})
337 b2 := bitmapToBytes(bits)
338 c.Assert(b2, qt.DeepEquals, b)
339
340 b = []byte{0, 15, 50}
341 bits = bytesToBitmap(b)
342 c.Assert(bits, qt.DeepEquals, []bool{false, false, false,
343 false, false, false, false, false, true, true, true, true,
344 false, false, false, false, false, true, false, false, true,
345 true, false, false})
346 b2 = bitmapToBytes(bits)
347 c.Assert(b2, qt.DeepEquals, b)
348
349 b = []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}
350 bits = bytesToBitmap(b)
351 b2 = bitmapToBytes(bits)
352 c.Assert(b2, qt.DeepEquals, b)
353
354 b = []byte("testbytes")
355 bits = bytesToBitmap(b)
356 b2 = bitmapToBytes(bits)
357 c.Assert(b2, qt.DeepEquals, b)
358}
359
360func TestPackAndUnpackSiblings(t *testing.T) {
361 c := qt.New(t)
362
363 siblingsHex := []string{
364 "0000000000000000000000000000000000000000000000000000000000000000",
365 "0100000000000000000000000000000000000000000000000000000000000000",
366 "0200000000000000000000000000000000000000000000000000000000000000",
367 "0000000000000000000000000000000000000000000000000000000000000000",
368 "0000000000000000000000000000000000000000000000000000000000000000",
369 "0000000000000000000000000000000000000000000000000000000000000000",
370 "0300000000000000000000000000000000000000000000000000000000000000",
371 "0400000000000000000000000000000000000000000000000000000000000000",
372 "0000000000000000000000000000000000000000000000000000000000000000",
373 "0000000000000000000000000000000000000000000000000000000000000000",
374 "0500000000000000000000000000000000000000000000000000000000000000",
375 }
376 siblings := make([][]byte, len(siblingsHex))
377 var err error
378 for i := 0; i < len(siblingsHex); i++ {
379 siblings[i], err = hex.DecodeString(siblingsHex[i])
380 c.Assert(err, qt.IsNil)
381 }
382
383 packed, err := PackSiblings(HashFunctionPoseidon, siblings)
384 c.Assert(err, qt.IsNil)
385 c.Assert(hex.EncodeToString(packed), qt.Equals, "a6000200c604"+
386 "0100000000000000000000000000000000000000000000000000000000000000"+
387 "0200000000000000000000000000000000000000000000000000000000000000"+
388 "0300000000000000000000000000000000000000000000000000000000000000"+
389 "0400000000000000000000000000000000000000000000000000000000000000"+
390 "0500000000000000000000000000000000000000000000000000000000000000")
391
392 unpacked, err := UnpackSiblings(HashFunctionPoseidon, packed)
393 c.Assert(err, qt.IsNil)
394 c.Assert(unpacked, qt.DeepEquals, siblings)
395
396 // another test with other values
397 siblingsHex = []string{
398 "1ce165cb1124ed3a0a94b4e212aaf7e8079f49b2fbef916bc290c593fda9092a",
399 "0000000000000000000000000000000000000000000000000000000000000000",
400 "0000000000000000000000000000000000000000000000000000000000000000",
401 "0000000000000000000000000000000000000000000000000000000000000000",
402 "33018202c57d898b84338b16d1a4960e133c6a4d656cfec1bd62a9ea00611729",
403 "bdbee2bd246ba0259a37be9a8740b550eed01c566aff0dca9a07306bcf731d13",
404 "0000000000000000000000000000000000000000000000000000000000000000",
405 "0000000000000000000000000000000000000000000000000000000000000000",
406 "d43b04d7c2d0bba83b4291fea9ba0fec7830d17af54cbe9967fe90b8244d4e0d",
407 "0000000000000000000000000000000000000000000000000000000000000000",
408 "7def274dbb3a72dca44f01a8d9f2f21a5be84c171eecef8e2e4112e7277e262a",
409 }
410 siblings = make([][]byte, len(siblingsHex))
411 for i := 0; i < len(siblingsHex); i++ {
412 siblings[i], err = hex.DecodeString(siblingsHex[i])
413 c.Assert(err, qt.IsNil)
414 }
415
416 packed, err = PackSiblings(HashFunctionPoseidon, siblings)
417 c.Assert(err, qt.IsNil)
418 c.Assert(hex.EncodeToString(packed), qt.Equals, "a60002003105"+
419 "1ce165cb1124ed3a0a94b4e212aaf7e8079f49b2fbef916bc290c593fda9092a"+
420 "33018202c57d898b84338b16d1a4960e133c6a4d656cfec1bd62a9ea00611729"+
421 "bdbee2bd246ba0259a37be9a8740b550eed01c566aff0dca9a07306bcf731d13"+
422 "d43b04d7c2d0bba83b4291fea9ba0fec7830d17af54cbe9967fe90b8244d4e0d"+
423 "7def274dbb3a72dca44f01a8d9f2f21a5be84c171eecef8e2e4112e7277e262a")
424
425 unpacked, err = UnpackSiblings(HashFunctionPoseidon, packed)
426 c.Assert(err, qt.IsNil)
427 c.Assert(unpacked, qt.DeepEquals, siblings)
428}
429
430func TestGenProofAndVerify(t *testing.T) {
431 c := qt.New(t)
432 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
433 c.Assert(err, qt.IsNil)
434 tree, err := NewTree(Config{Database: database, MaxLevels: 256,
435 HashFunction: HashFunctionPoseidon})
436 c.Assert(err, qt.IsNil)
437 defer tree.db.Close() //nolint:errcheck
438
439 bLen := 32
440 for i := 0; i < 10; i++ {
441 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
442 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
443 if err := tree.Add(k, v); err != nil {
444 t.Fatal(err)
445 }
446 }
447
448 k := BigIntToBytes(bLen, big.NewInt(int64(7)))
449 v := BigIntToBytes(bLen, big.NewInt(int64(14)))
450 kAux, proofV, siblings, existence, err := tree.GenProof(k)
451 c.Assert(err, qt.IsNil)
452 c.Assert(proofV, qt.DeepEquals, v)
453 c.Assert(k, qt.DeepEquals, kAux)
454 c.Assert(existence, qt.IsTrue)
455
456 root, err := tree.Root()
457 c.Assert(err, qt.IsNil)
458 verif, err := CheckProof(tree.hashFunction, k, v, root, siblings)
459 c.Assert(err, qt.IsNil)
460 c.Check(verif, qt.IsTrue)
461}
462
463func TestDumpAndImportDump(t *testing.T) {
464 testDumpAndImportDump(t, false)
465}
466
467func TestDumpAndImportDumpInFile(t *testing.T) {
468 testDumpAndImportDump(t, true)
469}
470
471func testDumpAndImportDump(t *testing.T, inFile bool) {
472 c := qt.New(t)
473 database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
474 c.Assert(err, qt.IsNil)
475 tree1, err := NewTree(Config{Database: database1, MaxLevels: 256,
476 HashFunction: HashFunctionPoseidon})
477 c.Assert(err, qt.IsNil)
478 defer tree1.db.Close() //nolint:errcheck
479
480 bLen := 32
481 for i := 0; i < 16; i++ {
482 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
483 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
484 if err := tree1.Add(k, v); err != nil {
485 t.Fatal(err)
486 }
487 }
488
489 var e []byte
490 filePath := c.TempDir()
491 fileName := filepath.Join(filePath, "dump.bin")
492 if inFile {
493 f, err := os.Create(fileName)
494 c.Assert(err, qt.IsNil)
495 w := bufio.NewWriter(f)
496 err = tree1.DumpWriter(nil, w)
497 c.Assert(err, qt.IsNil)
498 } else {
499 e, err = tree1.Dump(nil)
500 c.Assert(err, qt.IsNil)
501 }
502
503 database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
504 c.Assert(err, qt.IsNil)
505 tree2, err := NewTree(Config{Database: database2, MaxLevels: 256,
506 HashFunction: HashFunctionPoseidon})
507 c.Assert(err, qt.IsNil)
508 defer tree2.db.Close() //nolint:errcheck
509
510 if inFile {
511 f, err := os.Open(filepath.Clean(fileName))
512 c.Assert(err, qt.IsNil)
513 r := bufio.NewReader(f)
514 err = tree2.ImportDumpReader(r)
515 c.Assert(err, qt.IsNil)
516 } else {
517 err = tree2.ImportDump(e)
518 c.Assert(err, qt.IsNil)
519 }
520
521 root1, err := tree1.Root()
522 c.Assert(err, qt.IsNil)
523 root2, err := tree2.Root()
524 c.Assert(err, qt.IsNil)
525 c.Check(root2, qt.DeepEquals, root1)
526 c.Check(hex.EncodeToString(root2), qt.Equals,
527 "0d93aaa3362b2f999f15e15728f123087c2eee716f01c01f56e23aae07f09f08")
528}
529
530func TestRWMutex(t *testing.T) {
531 c := qt.New(t)
532 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
533 c.Assert(err, qt.IsNil)
534 tree, err := NewTree(Config{Database: database, MaxLevels: 256,
535 HashFunction: HashFunctionPoseidon})
536 c.Assert(err, qt.IsNil)
537 defer tree.db.Close() //nolint:errcheck
538
539 bLen := 32
540 var keys, values [][]byte
541 for i := 0; i < 1000; i++ {
542 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
543 v := BigIntToBytes(bLen, big.NewInt(0))
544 keys = append(keys, k)
545 values = append(values, v)
546 }
547 go func() {
548 _, err = tree.AddBatch(keys, values)
549 if err != nil {
550 panic(err)
551 }
552 }()
553
554 time.Sleep(500 * time.Millisecond)
555 k := BigIntToBytes(bLen, big.NewInt(int64(99999)))
556 v := BigIntToBytes(bLen, big.NewInt(int64(99999)))
557 if err := tree.Add(k, v); err != nil {
558 t.Fatal(err)
559 }
560}
561
562// TODO UPDATE
563// func TestSetGetNLeafs(t *testing.T) {
564// c := qt.New(t)
565// database, err := badgerdb.New(db.Options{Path: c.TempDir()})
566// c.Assert(err, qt.IsNil)
567// tree, err := NewTree(database, 100, HashFunctionPoseidon)
568// c.Assert(err, qt.IsNil)
569//
570// // 0
571// tree.dbBatch = tree.db.NewBatch()
572//
573// err = tree.setNLeafs(0)
574// c.Assert(err, qt.IsNil)
575//
576// err = tree.dbBatch.Write()
577// c.Assert(err, qt.IsNil)
578//
579// n, err := tree.GetNLeafs()
580// c.Assert(err, qt.IsNil)
581// c.Assert(n, qt.Equals, 0)
582//
583// // 1024
584// tree.dbBatch = tree.db.NewBatch()
585//
586// err = tree.setNLeafs(1024)
587// c.Assert(err, qt.IsNil)
588//
589// err = tree.dbBatch.Write()
590// c.Assert(err, qt.IsNil)
591//
592// n, err = tree.GetNLeafs()
593// c.Assert(err, qt.IsNil)
594// c.Assert(n, qt.Equals, 1024)
595//
596// // 2**64 -1
597// tree.dbBatch = tree.db.NewBatch()
598//
599// maxUint := ^uint(0)
600// maxInt := int(maxUint >> 1)
601//
602// err = tree.setNLeafs(maxInt)
603// c.Assert(err, qt.IsNil)
604//
605// err = tree.dbBatch.Write()
606// c.Assert(err, qt.IsNil)
607//
608// n, err = tree.GetNLeafs()
609// c.Assert(err, qt.IsNil)
610// c.Assert(n, qt.Equals, maxInt)
611// }
612
613func TestAddBatchFullyUsed(t *testing.T) {
614 c := qt.New(t)
615
616 database1, err := badgerdb.New(db.Options{Path: c.TempDir()})
617 c.Assert(err, qt.IsNil)
618 tree1, err := NewTree(Config{Database: database1, MaxLevels: 4,
619 HashFunction: HashFunctionPoseidon})
620 c.Assert(err, qt.IsNil)
621
622 database2, err := badgerdb.New(db.Options{Path: c.TempDir()})
623 c.Assert(err, qt.IsNil)
624 tree2, err := NewTree(Config{Database: database2, MaxLevels: 4,
625 HashFunction: HashFunctionPoseidon})
626 c.Assert(err, qt.IsNil)
627
628 var keys, values [][]byte
629 for i := 0; i < 16; i++ {
630 k := BigIntToBytes(1, big.NewInt(int64(i)))
631 v := k
632
633 keys = append(keys, k)
634 values = append(values, v)
635
636 // add one by one expecting no error
637 err := tree1.Add(k, v)
638 c.Assert(err, qt.IsNil)
639 }
640
641 invalids, err := tree2.AddBatch(keys, values)
642 c.Assert(err, qt.IsNil)
643 c.Assert(0, qt.Equals, len(invalids))
644
645 root1, err := tree1.Root()
646 c.Assert(err, qt.IsNil)
647 root2, err := tree2.Root()
648 c.Assert(err, qt.IsNil)
649 c.Assert(root1, qt.DeepEquals, root2)
650
651 // get all key-values and check that are equal between both trees
652 for i := 0; i < 16; i++ {
653 auxK1, auxV1, err := tree1.Get(BigIntToBytes(1, big.NewInt(int64(i))))
654 c.Assert(err, qt.IsNil)
655
656 auxK2, auxV2, err := tree2.Get(BigIntToBytes(1, big.NewInt(int64(i))))
657 c.Assert(err, qt.IsNil)
658
659 c.Assert(auxK1, qt.DeepEquals, auxK2)
660 c.Assert(auxV1, qt.DeepEquals, auxV2)
661 }
662
663 // try adding one more key to both trees (through Add & AddBatch) and
664 // expect not being added due the tree is already full
665 k := BigIntToBytes(1, big.NewInt(int64(16)))
666 v := k
667 err = tree1.Add(k, v)
668 c.Assert(err, qt.Equals, ErrMaxVirtualLevel)
669
670 invalids, err = tree2.AddBatch([][]byte{k}, [][]byte{v})
671 c.Assert(err, qt.IsNil)
672 c.Assert(1, qt.Equals, len(invalids))
673}
674
675func TestSetRoot(t *testing.T) {
676 c := qt.New(t)
677 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
678 c.Assert(err, qt.IsNil)
679 tree, err := NewTree(Config{Database: database, MaxLevels: 256,
680 HashFunction: HashFunctionPoseidon})
681 c.Assert(err, qt.IsNil)
682
683 expectedRoot := "13742386369878513332697380582061714160370929283209286127733983161245560237407"
684
685 // fill the tree
686 bLen := 32
687 var keys, values [][]byte
688 for i := 0; i < 1000; i++ {
689 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
690 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
691 keys = append(keys, k)
692 values = append(values, v)
693 }
694 indexes, err := tree.AddBatch(keys, values)
695 c.Assert(err, qt.IsNil)
696 c.Check(len(indexes), qt.Equals, 0)
697 checkRootBIString(c, tree,
698 expectedRoot)
699
700 // add one more k-v
701 k := BigIntToBytes(bLen, big.NewInt(1000))
702 v := BigIntToBytes(bLen, big.NewInt(1000))
703 err = tree.Add(k, v)
704 c.Assert(err, qt.IsNil)
705 checkRootBIString(c, tree,
706 "10747149055773881257049574592162159501044114324358186833013814735296193179713")
707
708 // do a SetRoot, and expect the same root than the original tree
709 pastRootBI, ok := new(big.Int).SetString(expectedRoot, 10)
710 c.Assert(ok, qt.IsTrue)
711 pastRoot := BigIntToBytes(32, pastRootBI)
712
713 err = tree.SetRoot(pastRoot)
714 c.Assert(err, qt.IsNil)
715 checkRootBIString(c, tree, expectedRoot)
716
717 // check that the tree can be updated
718 err = tree.Add([]byte("test"), []byte("test"))
719 c.Assert(err, qt.IsNil)
720 err = tree.Update([]byte("test"), []byte("test"))
721 c.Assert(err, qt.IsNil)
722
723 // check that the k-v '1000' does not exist in the new tree
724 _, _, err = tree.Get(k)
725 c.Assert(err, qt.Equals, ErrKeyNotFound)
726
727 // check that can be set an empty root
728 err = tree.SetRoot(tree.emptyHash)
729 c.Assert(err, qt.IsNil)
730}
731
732func TestSnapshot(t *testing.T) {
733 c := qt.New(t)
734 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
735 c.Assert(err, qt.IsNil)
736 tree, err := NewTree(Config{Database: database, MaxLevels: 256,
737 HashFunction: HashFunctionPoseidon})
738 c.Assert(err, qt.IsNil)
739
740 // fill the tree
741 bLen := 32
742 var keys, values [][]byte
743 for i := 0; i < 1000; i++ {
744 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
745 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
746 keys = append(keys, k)
747 values = append(values, v)
748 }
749 indexes, err := tree.AddBatch(keys, values)
750 c.Assert(err, qt.IsNil)
751 c.Check(len(indexes), qt.Equals, 0)
752 checkRootBIString(c, tree,
753 "13742386369878513332697380582061714160370929283209286127733983161245560237407")
754
755 // do a snapshot, and expect the same root than the original tree
756 snapshotTree, err := tree.Snapshot(nil)
757 c.Assert(err, qt.IsNil)
758 checkRootBIString(c, snapshotTree,
759 "13742386369878513332697380582061714160370929283209286127733983161245560237407")
760
761 // check that the snapshotTree can not be updated
762 _, err = snapshotTree.AddBatch(keys, values)
763 c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
764 err = snapshotTree.Add([]byte("test"), []byte("test"))
765 c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
766 err = snapshotTree.Update([]byte("test"), []byte("test"))
767 c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
768 err = snapshotTree.ImportDump(nil)
769 c.Assert(err, qt.Equals, ErrSnapshotNotEditable)
770
771 // update the original tree by adding a new key-value, and check that
772 // snapshotTree still has the old root, but the original tree has a new
773 // root
774 err = tree.Add([]byte("test"), []byte("test"))
775 c.Assert(err, qt.IsNil)
776 checkRootBIString(c, snapshotTree,
777 "13742386369878513332697380582061714160370929283209286127733983161245560237407")
778 checkRootBIString(c, tree,
779 "1025190963769001718196479367844646783678188389989148142691917685159698888868")
780}
781
782func TestGetFromSnapshotExpectArboErrKeyNotFound(t *testing.T) {
783 c := qt.New(t)
784
785 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
786 c.Assert(err, qt.IsNil)
787 tree, err := NewTree(Config{Database: database, MaxLevels: 256,
788 HashFunction: HashFunctionPoseidon})
789 c.Assert(err, qt.IsNil)
790 defer tree.db.Close() //nolint:errcheck
791
792 bLen := 32
793 k := BigIntToBytes(bLen, big.NewInt(int64(3)))
794
795 root, err := tree.Root()
796 c.Assert(err, qt.IsNil)
797 tree, err = tree.Snapshot(root)
798 c.Assert(err, qt.IsNil)
799
800 _, _, err = tree.Get(k)
801 c.Assert(err, qt.Equals, ErrKeyNotFound) // and not equal to db.ErrKeyNotFound
802}
803
804func TestKeyLen(t *testing.T) {
805 c := qt.New(t)
806 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
807 c.Assert(err, qt.IsNil)
808 // maxLevels is 100, keyPath length = ceil(maxLevels/8) = 13
809 maxLevels := 100
810 tree, err := NewTree(Config{Database: database, MaxLevels: maxLevels,
811 HashFunction: HashFunctionBlake2b})
812 c.Assert(err, qt.IsNil)
813
814 // expect no errors when adding a key of only 4 bytes (when the
815 // required length of keyPath for 100 levels would be 13 bytes)
816 bLen := 4
817 k := BigIntToBytes(bLen, big.NewInt(1))
818 v := BigIntToBytes(bLen, big.NewInt(1))
819
820 err = tree.Add(k, v)
821 c.Assert(err, qt.IsNil)
822
823 err = tree.Update(k, v)
824 c.Assert(err, qt.IsNil)
825
826 _, _, _, _, err = tree.GenProof(k)
827 c.Assert(err, qt.IsNil)
828
829 _, _, err = tree.Get(k)
830 c.Assert(err, qt.IsNil)
831
832 k = BigIntToBytes(bLen, big.NewInt(2))
833 v = BigIntToBytes(bLen, big.NewInt(2))
834 invalids, err := tree.AddBatch([][]byte{k}, [][]byte{v})
835 c.Assert(err, qt.IsNil)
836 c.Assert(len(invalids), qt.Equals, 0)
837
838 // expect errors when adding a key bigger than maximum capacity of the
839 // tree: ceil(maxLevels/8)
840 maxLevels = 32
841 database, err = badgerdb.New(db.Options{Path: c.TempDir()})
842 c.Assert(err, qt.IsNil)
843 tree, err = NewTree(Config{Database: database, MaxLevels: maxLevels,
844 HashFunction: HashFunctionBlake2b})
845 c.Assert(err, qt.IsNil)
846
847 maxKeyLen := int(math.Ceil(float64(maxLevels) / float64(8))) //nolint:gomnd
848 k = BigIntToBytes(maxKeyLen+1, big.NewInt(1))
849 v = BigIntToBytes(maxKeyLen+1, big.NewInt(1))
850
851 expectedErrMsg := "len(k) can not be bigger than ceil(maxLevels/8)," +
852 " where len(k): 5, maxLevels: 32, max key len=ceil(maxLevels/8): 4." +
853 " Might need a bigger tree depth (maxLevels>=40) in order to input" +
854 " keys of length 5"
855
856 err = tree.Add(k, v)
857 c.Assert(err.Error(), qt.Equals, expectedErrMsg)
858
859 err = tree.Update(k, v)
860 c.Assert(err.Error(), qt.Equals, expectedErrMsg)
861
862 _, _, _, _, err = tree.GenProof(k)
863 c.Assert(err.Error(), qt.Equals, expectedErrMsg)
864
865 _, _, err = tree.Get(k)
866 c.Assert(err.Error(), qt.Equals, expectedErrMsg)
867
868 // check AddBatch with few key-values
869 invalids, err = tree.AddBatch([][]byte{k}, [][]byte{v})
870 c.Assert(err, qt.IsNil)
871 c.Assert(len(invalids), qt.Equals, 1)
872
873 // check AddBatch with many key-values
874 nCPU := flp2(runtime.NumCPU())
875 nKVs := nCPU + 1
876 var ks, vs [][]byte
877 for i := 0; i < nKVs; i++ {
878 ks = append(ks, BigIntToBytes(maxKeyLen+1, big.NewInt(1)))
879 vs = append(vs, BigIntToBytes(maxKeyLen+1, big.NewInt(1)))
880 }
881 invalids, err = tree.AddBatch(ks, vs)
882 c.Assert(err, qt.IsNil)
883 c.Assert(len(invalids), qt.Equals, nKVs)
884
885 // check that with maxKeyLen it can be added
886 k = BigIntToBytes(maxKeyLen, big.NewInt(1))
887 err = tree.Add(k, v)
888 c.Assert(err, qt.IsNil)
889
890 // check CheckProof check with key longer
891 kAux, vAux, packedSiblings, existence, err := tree.GenProof(k)
892 c.Assert(err, qt.IsNil)
893 c.Assert(existence, qt.IsTrue)
894
895 root, err := tree.Root()
896 c.Assert(err, qt.IsNil)
897 verif, err := CheckProof(tree.HashFunction(), kAux, vAux, root, packedSiblings)
898 c.Assert(err, qt.IsNil)
899 c.Assert(verif, qt.IsTrue)
900
901 // use a similar key but with one zero, expect that CheckProof fails on
902 // the verification
903 kAux = append(kAux, 0)
904 verif, err = CheckProof(tree.HashFunction(), kAux, vAux, root, packedSiblings)
905 c.Assert(err, qt.IsNil)
906 c.Assert(verif, qt.IsFalse)
907}
908
909func TestKeyLenBiggerThan32(t *testing.T) {
910 c := qt.New(t)
911 maxLevels := 264
912 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
913 c.Assert(err, qt.IsNil)
914 tree, err := NewTree(Config{Database: database, MaxLevels: maxLevels,
915 HashFunction: HashFunctionBlake2b})
916 c.Assert(err, qt.IsNil)
917
918 bLen := 33
919 err = tree.Add(
920 randomBytes(bLen),
921 randomBytes(bLen))
922 c.Assert(err, qt.IsNil)
923
924 // 2nd key that we add, will find a node with len(key)==32 (due the
925 // hash output size, expect that next Add does not give any error, as
926 // internally it will use a keyPath of size corresponent to the
927 // maxLevels size of the tree
928 err = tree.Add(
929 randomBytes(bLen),
930 randomBytes(bLen))
931 c.Assert(err, qt.IsNil)
932}
933
934func BenchmarkAdd(b *testing.B) {
935 bLen := 32 // for both Poseidon & Sha256
936 // prepare inputs
937 var ks, vs [][]byte
938 for i := 0; i < 1000; i++ {
939 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
940 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
941 ks = append(ks, k)
942 vs = append(vs, v)
943 }
944
945 b.Run("Poseidon", func(b *testing.B) {
946 benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
947 })
948 b.Run("Sha256", func(b *testing.B) {
949 benchmarkAdd(b, HashFunctionSha256, ks, vs)
950 })
951}
952
953func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
954 c := qt.New(b)
955 database, err := badgerdb.New(db.Options{Path: c.TempDir()})
956 c.Assert(err, qt.IsNil)
957 tree, err := NewTree(Config{Database: database, MaxLevels: 140,
958 HashFunction: hashFunc})
959 c.Assert(err, qt.IsNil)
960 defer tree.db.Close() //nolint:errcheck
961
962 for i := 0; i < len(ks); i++ {
963 if err := tree.Add(ks[i], vs[i]); err != nil {
964 b.Fatal(err)
965 }
966 }
967}