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) { // TMP
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
224func TestGet(t *testing.T) {
225 c := qt.New(t)
226 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
227 c.Assert(err, qt.IsNil)
228 defer tree.db.Close()
229
230 for i := 0; i < 10; i++ {
231 k := BigIntToBytes(big.NewInt(int64(i)))
232 v := BigIntToBytes(big.NewInt(int64(i * 2)))
233 if err := tree.Add(k, v); err != nil {
234 t.Fatal(err)
235 }
236 }
237
238 k := BigIntToBytes(big.NewInt(int64(7)))
239 gettedKey, gettedValue, err := tree.Get(k)
240 c.Assert(err, qt.IsNil)
241 c.Check(gettedKey, qt.DeepEquals, k)
242 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(big.NewInt(int64(7*2))))
243}
244
245func TestGenProofAndVerify(t *testing.T) {
246 c := qt.New(t)
247 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
248 c.Assert(err, qt.IsNil)
249 defer tree.db.Close()
250
251 for i := 0; i < 10; i++ {
252 k := BigIntToBytes(big.NewInt(int64(i)))
253 v := BigIntToBytes(big.NewInt(int64(i * 2)))
254 if err := tree.Add(k, v); err != nil {
255 t.Fatal(err)
256 }
257 }
258
259 k := BigIntToBytes(big.NewInt(int64(7)))
260 siblings, err := tree.GenProof(k)
261 c.Assert(err, qt.IsNil)
262
263 k = BigIntToBytes(big.NewInt(int64(7)))
264 v := BigIntToBytes(big.NewInt(int64(14)))
265 verif, err := CheckProof(tree.hashFunction, k, v, tree.Root(), siblings)
266 c.Assert(err, qt.IsNil)
267 c.Check(verif, qt.IsTrue)
268}
269
270func TestDumpAndImportDump(t *testing.T) {
271 c := qt.New(t)
272 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
273 c.Assert(err, qt.IsNil)
274 defer tree1.db.Close()
275
276 for i := 0; i < 16; i++ {
277 k := BigIntToBytes(big.NewInt(int64(i)))
278 v := BigIntToBytes(big.NewInt(int64(i * 2)))
279 if err := tree1.Add(k, v); err != nil {
280 t.Fatal(err)
281 }
282 }
283
284 e, err := tree1.Dump()
285 c.Assert(err, qt.IsNil)
286
287 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
288 c.Assert(err, qt.IsNil)
289 defer tree2.db.Close()
290 err = tree2.ImportDump(e)
291 c.Assert(err, qt.IsNil)
292 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
293 c.Check(hex.EncodeToString(tree2.Root()), qt.Equals,
294 "0d93aaa3362b2f999f15e15728f123087c2eee716f01c01f56e23aae07f09f08")
295}
296
297func TestRWMutex(t *testing.T) {
298 c := qt.New(t)
299 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
300 c.Assert(err, qt.IsNil)
301 defer tree.db.Close()
302
303 var keys, values [][]byte
304 for i := 0; i < 1000; i++ {
305 k := BigIntToBytes(big.NewInt(int64(i)))
306 v := BigIntToBytes(big.NewInt(0))
307 keys = append(keys, k)
308 values = append(values, v)
309 }
310 go func() {
311 _, err = tree.AddBatch(keys, values)
312 if err != nil {
313 panic(err)
314 }
315 }()
316
317 time.Sleep(500 * time.Millisecond)
318 k := BigIntToBytes(big.NewInt(int64(99999)))
319 v := BigIntToBytes(big.NewInt(int64(99999)))
320 if err := tree.Add(k, v); err != nil {
321 t.Fatal(err)
322 }
323}
324
325func BenchmarkAdd(b *testing.B) {
326 // prepare inputs
327 var ks, vs [][]byte
328 for i := 0; i < 1000; i++ {
329 k := BigIntToBytes(big.NewInt(int64(i)))
330 v := BigIntToBytes(big.NewInt(int64(i)))
331 ks = append(ks, k)
332 vs = append(vs, v)
333 }
334
335 b.Run("Poseidon", func(b *testing.B) {
336 benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
337 })
338 b.Run("Sha256", func(b *testing.B) {
339 benchmarkAdd(b, HashFunctionSha256, ks, vs)
340 })
341}
342
343func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
344 c := qt.New(b)
345 tree, err := NewTree(memory.NewMemoryStorage(), 140, hashFunc)
346 c.Assert(err, qt.IsNil)
347 defer tree.db.Close()
348
349 for i := 0; i < len(ks); i++ {
350 if err := tree.Add(ks[i], vs[i]); err != nil {
351 b.Fatal(err)
352 }
353 }
354}