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 if err := tree.Add(k, v); err != nil {
145 t.Fatal(err)
146 }
147 err = tree.Add(k, v)
148 c.Assert(err, qt.Not(qt.IsNil))
149 c.Check(err, qt.Equals, ErrMaxVirtualLevel)
150}
151
152func TestUpdate(t *testing.T) {
153 c := qt.New(t)
154 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
155 c.Assert(err, qt.IsNil)
156 defer tree.db.Close()
157
158 bLen := tree.HashFunction().Len()
159 k := BigIntToBytes(bLen, big.NewInt(int64(20)))
160 v := BigIntToBytes(bLen, big.NewInt(int64(12)))
161 if err := tree.Add(k, v); err != nil {
162 t.Fatal(err)
163 }
164
165 v = BigIntToBytes(bLen, big.NewInt(int64(11)))
166 err = tree.Update(k, v)
167 c.Assert(err, qt.IsNil)
168
169 gettedKey, gettedValue, err := tree.Get(k)
170 c.Assert(err, qt.IsNil)
171 c.Check(gettedKey, qt.DeepEquals, k)
172 c.Check(gettedValue, qt.DeepEquals, v)
173
174 // add more leafs to the tree to do another test
175 for i := 0; i < 16; i++ {
176 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
177 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
178 if err := tree.Add(k, v); err != nil {
179 t.Fatal(err)
180 }
181 }
182
183 k = BigIntToBytes(bLen, big.NewInt(int64(3)))
184 v = BigIntToBytes(bLen, big.NewInt(int64(11)))
185 // check that before the Update, value for 3 is !=11
186 gettedKey, gettedValue, err = tree.Get(k)
187 c.Assert(err, qt.IsNil)
188 c.Check(gettedKey, qt.DeepEquals, k)
189 c.Check(gettedValue, qt.Not(qt.DeepEquals), v)
190 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(6)))
191
192 err = tree.Update(k, v)
193 c.Assert(err, qt.IsNil)
194
195 // check that after Update, the value for 3 is ==11
196 gettedKey, gettedValue, err = tree.Get(k)
197 c.Assert(err, qt.IsNil)
198 c.Check(gettedKey, qt.DeepEquals, k)
199 c.Check(gettedValue, qt.DeepEquals, v)
200 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(11)))
201}
202
203func TestAux(t *testing.T) { // TODO split in proper tests
204 c := qt.New(t)
205 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
206 c.Assert(err, qt.IsNil)
207 defer tree.db.Close()
208
209 bLen := tree.HashFunction().Len()
210 k := BigIntToBytes(bLen, big.NewInt(int64(1)))
211 v := BigIntToBytes(bLen, big.NewInt(int64(0)))
212 err = tree.Add(k, v)
213 c.Assert(err, qt.IsNil)
214 k = BigIntToBytes(bLen, big.NewInt(int64(256)))
215 err = tree.Add(k, v)
216 c.Assert(err, qt.IsNil)
217
218 k = BigIntToBytes(bLen, big.NewInt(int64(257)))
219 err = tree.Add(k, v)
220 c.Assert(err, qt.IsNil)
221
222 k = BigIntToBytes(bLen, big.NewInt(int64(515)))
223 err = tree.Add(k, v)
224 c.Assert(err, qt.IsNil)
225 k = BigIntToBytes(bLen, big.NewInt(int64(770)))
226 err = tree.Add(k, v)
227 c.Assert(err, qt.IsNil)
228
229 k = BigIntToBytes(bLen, big.NewInt(int64(388)))
230 err = tree.Add(k, v)
231 c.Assert(err, qt.IsNil)
232 k = BigIntToBytes(bLen, big.NewInt(int64(900)))
233 err = tree.Add(k, v)
234 c.Assert(err, qt.IsNil)
235 //
236 // err = tree.PrintGraphviz(nil)
237 // c.Assert(err, qt.IsNil)
238}
239
240func TestGet(t *testing.T) {
241 c := qt.New(t)
242 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
243 c.Assert(err, qt.IsNil)
244 defer tree.db.Close()
245
246 bLen := tree.HashFunction().Len()
247 for i := 0; i < 10; i++ {
248 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
249 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
250 if err := tree.Add(k, v); err != nil {
251 t.Fatal(err)
252 }
253 }
254
255 k := BigIntToBytes(bLen, big.NewInt(int64(7)))
256 gettedKey, gettedValue, err := tree.Get(k)
257 c.Assert(err, qt.IsNil)
258 c.Check(gettedKey, qt.DeepEquals, k)
259 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(bLen, big.NewInt(int64(7*2))))
260}
261
262func TestGenProofAndVerify(t *testing.T) {
263 c := qt.New(t)
264 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
265 c.Assert(err, qt.IsNil)
266 defer tree.db.Close()
267
268 bLen := tree.HashFunction().Len()
269 for i := 0; i < 10; i++ {
270 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
271 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
272 if err := tree.Add(k, v); err != nil {
273 t.Fatal(err)
274 }
275 }
276
277 k := BigIntToBytes(bLen, big.NewInt(int64(7)))
278 v := BigIntToBytes(bLen, big.NewInt(int64(14)))
279 proofV, siblings, err := tree.GenProof(k)
280 c.Assert(err, qt.IsNil)
281 c.Assert(proofV, qt.DeepEquals, v)
282
283 verif, err := CheckProof(tree.hashFunction, k, v, tree.Root(), siblings)
284 c.Assert(err, qt.IsNil)
285 c.Check(verif, qt.IsTrue)
286}
287
288func TestDumpAndImportDump(t *testing.T) {
289 c := qt.New(t)
290 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
291 c.Assert(err, qt.IsNil)
292 defer tree1.db.Close()
293
294 bLen := tree1.HashFunction().Len()
295 for i := 0; i < 16; i++ {
296 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
297 v := BigIntToBytes(bLen, big.NewInt(int64(i*2)))
298 if err := tree1.Add(k, v); err != nil {
299 t.Fatal(err)
300 }
301 }
302
303 e, err := tree1.Dump(nil)
304 c.Assert(err, qt.IsNil)
305
306 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
307 c.Assert(err, qt.IsNil)
308 defer tree2.db.Close()
309 err = tree2.ImportDump(e)
310 c.Assert(err, qt.IsNil)
311 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
312 c.Check(hex.EncodeToString(tree2.Root()), qt.Equals,
313 "0d93aaa3362b2f999f15e15728f123087c2eee716f01c01f56e23aae07f09f08")
314}
315
316func TestRWMutex(t *testing.T) {
317 c := qt.New(t)
318 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
319 c.Assert(err, qt.IsNil)
320 defer tree.db.Close()
321
322 bLen := tree.HashFunction().Len()
323 var keys, values [][]byte
324 for i := 0; i < 1000; i++ {
325 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
326 v := BigIntToBytes(bLen, big.NewInt(0))
327 keys = append(keys, k)
328 values = append(values, v)
329 }
330 go func() {
331 _, err = tree.AddBatch(keys, values)
332 if err != nil {
333 panic(err)
334 }
335 }()
336
337 time.Sleep(500 * time.Millisecond)
338 k := BigIntToBytes(bLen, big.NewInt(int64(99999)))
339 v := BigIntToBytes(bLen, big.NewInt(int64(99999)))
340 if err := tree.Add(k, v); err != nil {
341 t.Fatal(err)
342 }
343}
344
345func TestSetGetNLeafs(t *testing.T) {
346 c := qt.New(t)
347 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
348 c.Assert(err, qt.IsNil)
349
350 // 0
351 tree.tx, err = tree.db.NewTx()
352 c.Assert(err, qt.IsNil)
353
354 err = tree.setNLeafs(0)
355 c.Assert(err, qt.IsNil)
356
357 err = tree.tx.Commit()
358 c.Assert(err, qt.IsNil)
359
360 n, err := tree.GetNLeafs()
361 c.Assert(err, qt.IsNil)
362 c.Assert(n, qt.Equals, 0)
363
364 // 1024
365 tree.tx, err = tree.db.NewTx()
366 c.Assert(err, qt.IsNil)
367
368 err = tree.setNLeafs(1024)
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, 1024)
377
378 // 2**64 -1
379 tree.tx, err = tree.db.NewTx()
380 c.Assert(err, qt.IsNil)
381
382 maxUint := ^uint(0)
383 maxInt := int(maxUint >> 1)
384
385 err = tree.setNLeafs(maxInt)
386 c.Assert(err, qt.IsNil)
387
388 err = tree.tx.Commit()
389 c.Assert(err, qt.IsNil)
390
391 n, err = tree.GetNLeafs()
392 c.Assert(err, qt.IsNil)
393 c.Assert(n, qt.Equals, maxInt)
394}
395
396func BenchmarkAdd(b *testing.B) {
397 bLen := 32 // for both Poseidon & Sha256
398 // prepare inputs
399 var ks, vs [][]byte
400 for i := 0; i < 1000; i++ {
401 k := BigIntToBytes(bLen, big.NewInt(int64(i)))
402 v := BigIntToBytes(bLen, big.NewInt(int64(i)))
403 ks = append(ks, k)
404 vs = append(vs, v)
405 }
406
407 b.Run("Poseidon", func(b *testing.B) {
408 benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
409 })
410 b.Run("Sha256", func(b *testing.B) {
411 benchmarkAdd(b, HashFunctionSha256, ks, vs)
412 })
413}
414
415func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
416 c := qt.New(b)
417 tree, err := NewTree(memory.NewMemoryStorage(), 140, hashFunc)
418 c.Assert(err, qt.IsNil)
419 defer tree.db.Close()
420
421 for i := 0; i < len(ks); i++ {
422 if err := tree.Add(ks[i], vs[i]); err != nil {
423 b.Fatal(err)
424 }
425 }
426}