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