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