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