arbo/tree_test.go
tree_test.goRaw
1package arbo
2
3import (
4 "encoding/hex"
5 "math/big"
6 "testing"
7 "time"
8
9 qt "github.com/frankban/quicktest"
10 "github.com/iden3/go-merkletree/db/memory"
11)
12
13func TestAddTestVectors(t *testing.T) {
14 c := qt.New(t)
15
16 // Poseidon test vectors generated using https://github.com/iden3/circomlib smt.js
17 testVectorsPoseidon := []string{
18 "0000000000000000000000000000000000000000000000000000000000000000",
19 "13578938674299138072471463694055224830892726234048532520316387704878000008795",
20 "5412393676474193513566895793055462193090331607895808993925969873307089394741",
21 "14204494359367183802864593755198662203838502594566452929175967972147978322084",
22 }
23 testAdd(c, HashFunctionPoseidon, testVectorsPoseidon)
24
25 testVectorsSha256 := []string{
26 "0000000000000000000000000000000000000000000000000000000000000000",
27 "46910109172468462938850740851377282682950237270676610513794735904325820156367",
28 "59481735341404520835410489183267411392292882901306595567679529387376287440550",
29 "20573794434149960984975763118181266662429997821552560184909083010514790081771",
30 }
31 testAdd(c, HashFunctionSha256, testVectorsSha256)
32}
33
34func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
35 tree, err := NewTree(memory.NewMemoryStorage(), 10, hashFunc)
36 c.Assert(err, qt.IsNil)
37 defer tree.db.Close()
38
39 c.Check(hex.EncodeToString(tree.Root()), qt.Equals, testVectors[0])
40
41 bLen := hashFunc.Len()
42 err = tree.Add(
43 BigIntToBytes(bLen, big.NewInt(1)),
44 BigIntToBytes(bLen, big.NewInt(2)))
45 c.Assert(err, qt.IsNil)
46 rootBI := BytesToBigInt(tree.Root())
47 c.Check(rootBI.String(), qt.Equals, testVectors[1])
48
49 err = tree.Add(
50 BigIntToBytes(bLen, big.NewInt(33)),
51 BigIntToBytes(bLen, big.NewInt(44)))
52 c.Assert(err, qt.IsNil)
53 rootBI = BytesToBigInt(tree.Root())
54 c.Check(rootBI.String(), qt.Equals, testVectors[2])
55
56 err = tree.Add(
57 BigIntToBytes(bLen, big.NewInt(1234)),
58 BigIntToBytes(bLen, big.NewInt(9876)))
59 c.Assert(err, qt.IsNil)
60 rootBI = BytesToBigInt(tree.Root())
61 c.Check(rootBI.String(), qt.Equals, testVectors[3])
62}
63
64func TestAddBatch(t *testing.T) {
65 c := qt.New(t)
66 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
67 c.Assert(err, qt.IsNil)
68 defer tree.db.Close()
69
70 bLen := tree.HashFunction().Len()
71 for i := 0; i < 1000; i++ {
72 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
73 v := BigIntToBytes(bLen, big.NewInt(0))
74 if err := tree.Add(k, v); err != nil {
75 t.Fatal(err)
76 }
77 }
78
79 rootBI := BytesToBigInt(tree.Root())
80 c.Check(rootBI.String(), qt.Equals,
81 "296519252211642170490407814696803112091039265640052570497930797516015811235")
82
83 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
84 c.Assert(err, qt.IsNil)
85 defer tree2.db.Close()
86
87 var keys, values [][]byte
88 for i := 0; i < 1000; i++ {
89 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
90 v := BigIntToBytes(bLen, big.NewInt(0))
91 keys = append(keys, k)
92 values = append(values, v)
93 }
94 indexes, err := tree2.AddBatch(keys, values)
95 c.Assert(err, qt.IsNil)
96 c.Check(len(indexes), qt.Equals, 0)
97
98 rootBI = BytesToBigInt(tree2.Root())
99 c.Check(rootBI.String(), qt.Equals,
100 "296519252211642170490407814696803112091039265640052570497930797516015811235")
101}
102
103func TestAddDifferentOrder(t *testing.T) {
104 c := qt.New(t)
105 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
106 c.Assert(err, qt.IsNil)
107 defer tree1.db.Close()
108
109 bLen := tree1.HashFunction().Len()
110 for i := 0; i < 16; i++ {
111 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
112 v := BigIntToBytes(bLen, big.NewInt(0))
113 if err := tree1.Add(k, v); err != nil {
114 t.Fatal(err)
115 }
116 }
117
118 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
119 c.Assert(err, qt.IsNil)
120 defer tree2.db.Close()
121
122 for i := 16 - 1; i >= 0; i-- {
123 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
124 v := BigIntToBytes(bLen, big.NewInt(0))
125 if err := tree2.Add(k, v); err != nil {
126 t.Fatal(err)
127 }
128 }
129
130 c.Check(hex.EncodeToString(tree2.Root()), qt.Equals, hex.EncodeToString(tree1.Root()))
131 c.Check(hex.EncodeToString(tree1.Root()), qt.Equals,
132 "3b89100bec24da9275c87bc188740389e1d5accfc7d88ba5688d7fa96a00d82f")
133}
134
135func TestAddRepeatedIndex(t *testing.T) {
136 c := qt.New(t)
137 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
138 c.Assert(err, qt.IsNil)
139 defer tree.db.Close()
140
141 bLen := tree.HashFunction().Len()
142 k := BigIntToBytes(bLen, big.NewInt(int64(3)))
143 v := BigIntToBytes(bLen, big.NewInt(int64(12)))
144
145 err = tree.Add(k, v)
146 c.Assert(err, qt.IsNil)
147 err = tree.Add(k, v) // repeating same key-value
148 c.Check(err, qt.Equals, ErrKeyAlreadyExists)
149}
150
151func TestUpdate(t *testing.T) {
152 c := qt.New(t)
153 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
154 c.Assert(err, qt.IsNil)
155 defer tree.db.Close()
156
157 bLen := tree.HashFunction().Len()
158 k := BigIntToBytes(bLen, big.NewInt(int64(20)))
159 v := BigIntToBytes(bLen, big.NewInt(int64(12)))
160 if err := tree.Add(k, v); err != nil {
161 t.Fatal(err)
162 }
163
164 v = BigIntToBytes(bLen, big.NewInt(int64(11)))
165 err = tree.Update(k, v)
166 c.Assert(err, qt.IsNil)
167
168 gettedKey, gettedValue, err := tree.Get(k)
169 c.Assert(err, qt.IsNil)
170 c.Check(gettedKey, qt.DeepEquals, k)
171 c.Check(gettedValue, qt.DeepEquals, v)
172
173 // add more leafs to the tree to do another test
174 for i := 0; i < 16; i++ {
175 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
176 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
177 if err := tree.Add(k, v); err != nil {
178 t.Fatal(err)
179 }
180 }
181
182 k = BigIntToBytes(bLen, big.NewInt(int64(3)))
183 v = BigIntToBytes(bLen, big.NewInt(int64(11)))
184 // check that before the Update, value for 3 is !=11
185 gettedKey, gettedValue, err = tree.Get(k)
186 c.Assert(err, qt.IsNil)
187 c.Check(gettedKey, qt.DeepEquals, k)
188 c.Check(gettedValue, qt.Not(qt.DeepEquals), v)
189 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(6)))
190
191 err = tree.Update(k, v)
192 c.Assert(err, qt.IsNil)
193
194 // check that after Update, the value for 3 is ==11
195 gettedKey, gettedValue, err = tree.Get(k)
196 c.Assert(err, qt.IsNil)
197 c.Check(gettedKey, qt.DeepEquals, k)
198 c.Check(gettedValue, qt.DeepEquals, v)
199 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(11)))
200}
201
202func TestAux(t *testing.T) { // TODO split in proper tests
203 c := qt.New(t)
204 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
205 c.Assert(err, qt.IsNil)
206 defer tree.db.Close()
207
208 bLen := tree.HashFunction().Len()
209 k := BigIntToBytes(bLen, big.NewInt(int64(1)))
210 v := BigIntToBytes(bLen, big.NewInt(int64(0)))
211 err = tree.Add(k, v)
212 c.Assert(err, qt.IsNil)
213 k = BigIntToBytes(bLen, big.NewInt(int64(256)))
214 err = tree.Add(k, v)
215 c.Assert(err, qt.IsNil)
216
217 k = BigIntToBytes(bLen, big.NewInt(int64(257)))
218 err = tree.Add(k, v)
219 c.Assert(err, qt.IsNil)
220
221 k = BigIntToBytes(bLen, big.NewInt(int64(515)))
222 err = tree.Add(k, v)
223 c.Assert(err, qt.IsNil)
224 k = BigIntToBytes(bLen, big.NewInt(int64(770)))
225 err = tree.Add(k, v)
226 c.Assert(err, qt.IsNil)
227
228 k = BigIntToBytes(bLen, big.NewInt(int64(388)))
229 err = tree.Add(k, v)
230 c.Assert(err, qt.IsNil)
231 k = BigIntToBytes(bLen, big.NewInt(int64(900)))
232 err = tree.Add(k, v)
233 c.Assert(err, qt.IsNil)
234 //
235 // err = tree.PrintGraphviz(nil)
236 // c.Assert(err, qt.IsNil)
237}
238
239func TestGet(t *testing.T) {
240 c := qt.New(t)
241 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
242 c.Assert(err, qt.IsNil)
243 defer tree.db.Close()
244
245 bLen := tree.HashFunction().Len()
246 for i := 0; i < 10; i++ {
247 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
248 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
249 if err := tree.Add(k, v); err != nil {
250 t.Fatal(err)
251 }
252 }
253
254 k := BigIntToBytes(bLen, big.NewInt(int64(7)))
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, BigIntToBytes(bLen, big.NewInt(int64(7*2))))
259}
260
261func TestGenProofAndVerify(t *testing.T) {
262 c := qt.New(t)
263 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
264 c.Assert(err, qt.IsNil)
265 defer tree.db.Close()
266
267 bLen := tree.HashFunction().Len()
268 for i := 0; i < 10; i++ {
269 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
270 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
271 if err := tree.Add(k, v); err != nil {
272 t.Fatal(err)
273 }
274 }
275
276 k := BigIntToBytes(bLen, big.NewInt(int64(7)))
277 v := BigIntToBytes(bLen, big.NewInt(int64(14)))
278 proofV, siblings, err := tree.GenProof(k)
279 c.Assert(err, qt.IsNil)
280 c.Assert(proofV, qt.DeepEquals, v)
281
282 verif, err := CheckProof(tree.hashFunction, k, v, tree.Root(), siblings)
283 c.Assert(err, qt.IsNil)
284 c.Check(verif, qt.IsTrue)
285}
286
287func TestDumpAndImportDump(t *testing.T) {
288 c := qt.New(t)
289 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
290 c.Assert(err, qt.IsNil)
291 defer tree1.db.Close()
292
293 bLen := tree1.HashFunction().Len()
294 for i := 0; i < 16; i++ {
295 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
296 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
297 if err := tree1.Add(k, v); err != nil {
298 t.Fatal(err)
299 }
300 }
301
302 e, err := tree1.Dump(nil)
303 c.Assert(err, qt.IsNil)
304
305 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
306 c.Assert(err, qt.IsNil)
307 defer tree2.db.Close()
308 err = tree2.ImportDump(e)
309 c.Assert(err, qt.IsNil)
310 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
311 c.Check(hex.EncodeToString(tree2.Root()), qt.Equals,
312 "0d93aaa3362b2f999f15e15728f123087c2eee716f01c01f56e23aae07f09f08")
313}
314
315func TestRWMutex(t *testing.T) {
316 c := qt.New(t)
317 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
318 c.Assert(err, qt.IsNil)
319 defer tree.db.Close()
320
321 bLen := tree.HashFunction().Len()
322 var keys, values [][]byte
323 for i := 0; i < 1000; i++ {
324 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
325 v := BigIntToBytes(bLen, big.NewInt(0))
326 keys = append(keys, k)
327 values = append(values, v)
328 }
329 go func() {
330 _, err = tree.AddBatch(keys, values)
331 if err != nil {
332 panic(err)
333 }
334 }()
335
336 time.Sleep(500 * time.Millisecond)
337 k := BigIntToBytes(bLen, big.NewInt(int64(99999)))
338 v := BigIntToBytes(bLen, big.NewInt(int64(99999)))
339 if err := tree.Add(k, v); err != nil {
340 t.Fatal(err)
341 }
342}
343
344func TestSetGetNLeafs(t *testing.T) {
345 c := qt.New(t)
346 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
347 c.Assert(err, qt.IsNil)
348
349 // 0
350 tree.tx, err = tree.db.NewTx()
351 c.Assert(err, qt.IsNil)
352
353 err = tree.setNLeafs(0)
354 c.Assert(err, qt.IsNil)
355
356 err = tree.tx.Commit()
357 c.Assert(err, qt.IsNil)
358
359 n, err := tree.GetNLeafs()
360 c.Assert(err, qt.IsNil)
361 c.Assert(n, qt.Equals, 0)
362
363 // 1024
364 tree.tx, err = tree.db.NewTx()
365 c.Assert(err, qt.IsNil)
366
367 err = tree.setNLeafs(1024)
368 c.Assert(err, qt.IsNil)
369
370 err = tree.tx.Commit()
371 c.Assert(err, qt.IsNil)
372
373 n, err = tree.GetNLeafs()
374 c.Assert(err, qt.IsNil)
375 c.Assert(n, qt.Equals, 1024)
376
377 // 2**64 -1
378 tree.tx, err = tree.db.NewTx()
379 c.Assert(err, qt.IsNil)
380
381 maxUint := ^uint(0)
382 maxInt := int(maxUint >> 1)
383
384 err = tree.setNLeafs(maxInt)
385 c.Assert(err, qt.IsNil)
386
387 err = tree.tx.Commit()
388 c.Assert(err, qt.IsNil)
389
390 n, err = tree.GetNLeafs()
391 c.Assert(err, qt.IsNil)
392 c.Assert(n, qt.Equals, maxInt)
393}
394
395func BenchmarkAdd(b *testing.B) {
396 bLen := 32 // for both Poseidon & Sha256
397 // prepare inputs
398 var ks, vs [][]byte
399 for i := 0; i < 1000; i++ {
400 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
401 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
402 ks = append(ks, k)
403 vs = append(vs, v)
404 }
405
406 b.Run("Poseidon", func(b *testing.B) {
407 benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
408 })
409 b.Run("Sha256", func(b *testing.B) {
410 benchmarkAdd(b, HashFunctionSha256, ks, vs)
411 })
412}
413
414func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
415 c := qt.New(b)
416 tree, err := NewTree(memory.NewMemoryStorage(), 140, hashFunc)
417 c.Assert(err, qt.IsNil)
418 defer tree.db.Close()
419
420 for i := 0; i < len(ks); i++ {
421 if err := tree.Add(ks[i], vs[i]); err != nil {
422 b.Fatal(err)
423 }
424 }
425}