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 // err = tree.PrintGraphviz(nil)
224 // c.Assert(err, qt.IsNil)
225}
226
227func TestGet(t *testing.T) {
228 c := qt.New(t)
229 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
230 c.Assert(err, qt.IsNil)
231 defer tree.db.Close()
232
233 for i := 0; i < 10; i++ {
234 k := BigIntToBytes(big.NewInt(int64(i)))
235 v := BigIntToBytes(big.NewInt(int64(i * 2)))
236 if err := tree.Add(k, v); err != nil {
237 t.Fatal(err)
238 }
239 }
240
241 k := BigIntToBytes(big.NewInt(int64(7)))
242 gettedKey, gettedValue, err := tree.Get(k)
243 c.Assert(err, qt.IsNil)
244 c.Check(gettedKey, qt.DeepEquals, k)
245 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(big.NewInt(int64(7*2))))
246}
247
248func TestGenProofAndVerify(t *testing.T) {
249 c := qt.New(t)
250 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
251 c.Assert(err, qt.IsNil)
252 defer tree.db.Close()
253
254 for i := 0; i < 10; i++ {
255 k := BigIntToBytes(big.NewInt(int64(i)))
256 v := BigIntToBytes(big.NewInt(int64(i * 2)))
257 if err := tree.Add(k, v); err != nil {
258 t.Fatal(err)
259 }
260 }
261
262 k := BigIntToBytes(big.NewInt(int64(7)))
263 siblings, err := tree.GenProof(k)
264 c.Assert(err, qt.IsNil)
265
266 k = BigIntToBytes(big.NewInt(int64(7)))
267 v := BigIntToBytes(big.NewInt(int64(14)))
268 verif, err := CheckProof(tree.hashFunction, k, v, tree.Root(), siblings)
269 c.Assert(err, qt.IsNil)
270 c.Check(verif, qt.IsTrue)
271}
272
273func TestDumpAndImportDump(t *testing.T) {
274 c := qt.New(t)
275 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
276 c.Assert(err, qt.IsNil)
277 defer tree1.db.Close()
278
279 for i := 0; i < 16; i++ {
280 k := BigIntToBytes(big.NewInt(int64(i)))
281 v := BigIntToBytes(big.NewInt(int64(i * 2)))
282 if err := tree1.Add(k, v); err != nil {
283 t.Fatal(err)
284 }
285 }
286
287 e, err := tree1.Dump()
288 c.Assert(err, qt.IsNil)
289
290 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
291 c.Assert(err, qt.IsNil)
292 defer tree2.db.Close()
293 err = tree2.ImportDump(e)
294 c.Assert(err, qt.IsNil)
295 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
296 c.Check(hex.EncodeToString(tree2.Root()), qt.Equals,
297 "0d93aaa3362b2f999f15e15728f123087c2eee716f01c01f56e23aae07f09f08")
298}
299
300func TestRWMutex(t *testing.T) {
301 c := qt.New(t)
302 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
303 c.Assert(err, qt.IsNil)
304 defer tree.db.Close()
305
306 var keys, values [][]byte
307 for i := 0; i < 1000; i++ {
308 k := BigIntToBytes(big.NewInt(int64(i)))
309 v := BigIntToBytes(big.NewInt(0))
310 keys = append(keys, k)
311 values = append(values, v)
312 }
313 go func() {
314 _, err = tree.AddBatch(keys, values)
315 if err != nil {
316 panic(err)
317 }
318 }()
319
320 time.Sleep(500 * time.Millisecond)
321 k := BigIntToBytes(big.NewInt(int64(99999)))
322 v := BigIntToBytes(big.NewInt(int64(99999)))
323 if err := tree.Add(k, v); err != nil {
324 t.Fatal(err)
325 }
326}
327
328func TestSetGetNLeafs(t *testing.T) {
329 c := qt.New(t)
330 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
331 c.Assert(err, qt.IsNil)
332
333 // 0
334 tree.tx, err = tree.db.NewTx()
335 c.Assert(err, qt.IsNil)
336
337 err = tree.setNLeafs(0)
338 c.Assert(err, qt.IsNil)
339
340 err = tree.tx.Commit()
341 c.Assert(err, qt.IsNil)
342
343 n, err := tree.GetNLeafs()
344 c.Assert(err, qt.IsNil)
345 c.Assert(n, qt.Equals, 0)
346
347 // 1024
348 tree.tx, err = tree.db.NewTx()
349 c.Assert(err, qt.IsNil)
350
351 err = tree.setNLeafs(1024)
352 c.Assert(err, qt.IsNil)
353
354 err = tree.tx.Commit()
355 c.Assert(err, qt.IsNil)
356
357 n, err = tree.GetNLeafs()
358 c.Assert(err, qt.IsNil)
359 c.Assert(n, qt.Equals, 1024)
360
361 // 2**64 -1
362 tree.tx, err = tree.db.NewTx()
363 c.Assert(err, qt.IsNil)
364
365 maxUint := ^uint(0)
366 maxInt := int(maxUint >> 1)
367
368 err = tree.setNLeafs(maxInt)
369 c.Assert(err, qt.IsNil)
370
371 err = tree.tx.Commit()
372 c.Assert(err, qt.IsNil)
373
374 n, err = tree.GetNLeafs()
375 c.Assert(err, qt.IsNil)
376 c.Assert(n, qt.Equals, maxInt)
377}
378
379func BenchmarkAdd(b *testing.B) {
380 // prepare inputs
381 var ks, vs [][]byte
382 for i := 0; i < 1000; i++ {
383 k := BigIntToBytes(big.NewInt(int64(i)))
384 v := BigIntToBytes(big.NewInt(int64(i)))
385 ks = append(ks, k)
386 vs = append(vs, v)
387 }
388
389 b.Run("Poseidon", func(b *testing.B) {
390 benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
391 })
392 b.Run("Sha256", func(b *testing.B) {
393 benchmarkAdd(b, HashFunctionSha256, ks, vs)
394 })
395}
396
397func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
398 c := qt.New(b)
399 tree, err := NewTree(memory.NewMemoryStorage(), 140, hashFunc)
400 c.Assert(err, qt.IsNil)
401 defer tree.db.Close()
402
403 for i := 0; i < len(ks); i++ {
404 if err := tree.Add(ks[i], vs[i]); err != nil {
405 b.Fatal(err)
406 }
407 }
408}