arbo/tree_test.go
tree_test.goRaw
1package arbo
2
3import (
4 "encoding/hex"
5 "math/big"
6 "testing"
7
8 qt "github.com/frankban/quicktest"
9 "github.com/iden3/go-merkletree/db/memory"
10)
11
12func TestAddTestVectors(t *testing.T) {
13 c := qt.New(t)
14
15 // Poseidon test vectors generated using https://github.com/iden3/circomlib smt.js
16 testVectorsPoseidon := []string{
17 "0000000000000000000000000000000000000000000000000000000000000000",
18 "13578938674299138072471463694055224830892726234048532520316387704878000008795",
19 "5412393676474193513566895793055462193090331607895808993925969873307089394741",
20 "14204494359367183802864593755198662203838502594566452929175967972147978322084",
21 }
22 testAdd(c, HashFunctionPoseidon, testVectorsPoseidon)
23
24 testVectorsSha256 := []string{
25 "0000000000000000000000000000000000000000000000000000000000000000",
26 "46910109172468462938850740851377282682950237270676610513794735904325820156367",
27 "59481735341404520835410489183267411392292882901306595567679529387376287440550",
28 "20573794434149960984975763118181266662429997821552560184909083010514790081771",
29 }
30 testAdd(c, HashFunctionSha256, testVectorsSha256)
31}
32
33func testAdd(c *qt.C, hashFunc HashFunction, testVectors []string) {
34 tree, err := NewTree(memory.NewMemoryStorage(), 10, hashFunc)
35 c.Assert(err, qt.IsNil)
36 defer tree.db.Close()
37
38 c.Check(hex.EncodeToString(tree.Root()), qt.Equals, testVectors[0])
39
40 err = tree.Add(
41 BigIntToBytes(big.NewInt(1)),
42 BigIntToBytes(big.NewInt(2)))
43 c.Assert(err, qt.IsNil)
44 rootBI := BytesToBigInt(tree.Root())
45 c.Check(rootBI.String(), qt.Equals, testVectors[1])
46
47 err = tree.Add(
48 BigIntToBytes(big.NewInt(33)),
49 BigIntToBytes(big.NewInt(44)))
50 c.Assert(err, qt.IsNil)
51 rootBI = BytesToBigInt(tree.Root())
52 c.Check(rootBI.String(), qt.Equals, testVectors[2])
53
54 err = tree.Add(
55 BigIntToBytes(big.NewInt(1234)),
56 BigIntToBytes(big.NewInt(9876)))
57 c.Assert(err, qt.IsNil)
58 rootBI = BytesToBigInt(tree.Root())
59 c.Check(rootBI.String(), qt.Equals, testVectors[3])
60}
61
62func TestAddBatch(t *testing.T) {
63 c := qt.New(t)
64 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
65 c.Assert(err, qt.IsNil)
66 defer tree.db.Close()
67
68 for i := 0; i < 1000; i++ {
69 k := BigIntToBytes(big.NewInt(int64(i)))
70 v := BigIntToBytes(big.NewInt(0))
71 if err := tree.Add(k, v); err != nil {
72 t.Fatal(err)
73 }
74 }
75
76 rootBI := BytesToBigInt(tree.Root())
77 c.Check(rootBI.String(), qt.Equals,
78 "296519252211642170490407814696803112091039265640052570497930797516015811235")
79
80 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
81 c.Assert(err, qt.IsNil)
82 defer tree2.db.Close()
83
84 var keys, values [][]byte
85 for i := 0; i < 1000; i++ {
86 k := BigIntToBytes(big.NewInt(int64(i)))
87 v := BigIntToBytes(big.NewInt(0))
88 keys = append(keys, k)
89 values = append(values, v)
90 }
91 indexes, err := tree2.AddBatch(keys, values)
92 c.Assert(err, qt.IsNil)
93 c.Check(len(indexes), qt.Equals, 0)
94
95 rootBI = BytesToBigInt(tree2.Root())
96 c.Check(rootBI.String(), qt.Equals,
97 "296519252211642170490407814696803112091039265640052570497930797516015811235")
98}
99
100func TestAddDifferentOrder(t *testing.T) {
101 c := qt.New(t)
102 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
103 c.Assert(err, qt.IsNil)
104 defer tree1.db.Close()
105
106 for i := 0; i < 16; i++ {
107 k := BigIntToBytes(big.NewInt(int64(i)))
108 v := BigIntToBytes(big.NewInt(0))
109 if err := tree1.Add(k, v); err != nil {
110 t.Fatal(err)
111 }
112 }
113
114 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
115 c.Assert(err, qt.IsNil)
116 defer tree2.db.Close()
117
118 for i := 16 - 1; i >= 0; i-- {
119 k := BigIntToBytes(big.NewInt(int64(i)))
120 v := BigIntToBytes(big.NewInt(0))
121 if err := tree2.Add(k, v); err != nil {
122 t.Fatal(err)
123 }
124 }
125
126 c.Check(hex.EncodeToString(tree2.Root()), qt.Equals, hex.EncodeToString(tree1.Root()))
127 c.Check(hex.EncodeToString(tree1.Root()), qt.Equals,
128 "3b89100bec24da9275c87bc188740389e1d5accfc7d88ba5688d7fa96a00d82f")
129}
130
131func TestAddRepeatedIndex(t *testing.T) {
132 c := qt.New(t)
133 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
134 c.Assert(err, qt.IsNil)
135 defer tree.db.Close()
136
137 k := BigIntToBytes(big.NewInt(int64(3)))
138 v := BigIntToBytes(big.NewInt(int64(12)))
139 if err := tree.Add(k, v); err != nil {
140 t.Fatal(err)
141 }
142 err = tree.Add(k, v)
143 c.Assert(err, qt.Not(qt.IsNil))
144 c.Check(err, qt.ErrorMatches, "max virtual level 100")
145}
146
147func TestUpdate(t *testing.T) {
148 c := qt.New(t)
149 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
150 c.Assert(err, qt.IsNil)
151 defer tree.db.Close()
152
153 k := BigIntToBytes(big.NewInt(int64(20)))
154 v := BigIntToBytes(big.NewInt(int64(12)))
155 if err := tree.Add(k, v); err != nil {
156 t.Fatal(err)
157 }
158
159 v = BigIntToBytes(big.NewInt(int64(11)))
160 err = tree.Update(k, v)
161 c.Assert(err, qt.IsNil)
162
163 gettedKey, gettedValue, err := tree.Get(k)
164 c.Assert(err, qt.IsNil)
165 c.Check(gettedKey, qt.DeepEquals, k)
166 c.Check(gettedValue, qt.DeepEquals, v)
167
168 // add more leafs to the tree to do another test
169 for i := 0; i < 16; i++ {
170 k := BigIntToBytes(big.NewInt(int64(i)))
171 v := BigIntToBytes(big.NewInt(int64(i * 2)))
172 if err := tree.Add(k, v); err != nil {
173 t.Fatal(err)
174 }
175 }
176
177 k = BigIntToBytes(big.NewInt(int64(3)))
178 v = BigIntToBytes(big.NewInt(int64(11)))
179 // check that before the Update, value for 3 is !=11
180 gettedKey, gettedValue, err = tree.Get(k)
181 c.Assert(err, qt.IsNil)
182 c.Check(gettedKey, qt.DeepEquals, k)
183 c.Check(gettedValue, qt.Not(qt.DeepEquals), v)
184 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(big.NewInt(6)))
185
186 err = tree.Update(k, v)
187 c.Assert(err, qt.IsNil)
188
189 // check that after Update, the value for 3 is ==11
190 gettedKey, gettedValue, err = tree.Get(k)
191 c.Assert(err, qt.IsNil)
192 c.Check(gettedKey, qt.DeepEquals, k)
193 c.Check(gettedValue, qt.DeepEquals, v)
194 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(big.NewInt(11)))
195}
196
197func TestAux(t *testing.T) {
198 c := qt.New(t)
199 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
200 c.Assert(err, qt.IsNil)
201 defer tree.db.Close()
202
203 k := BigIntToBytes(big.NewInt(int64(1)))
204 v := BigIntToBytes(big.NewInt(int64(0)))
205 err = tree.Add(k, v)
206 c.Assert(err, qt.IsNil)
207 k = BigIntToBytes(big.NewInt(int64(256)))
208 err = tree.Add(k, v)
209 c.Assert(err, qt.IsNil)
210
211 k = BigIntToBytes(big.NewInt(int64(257)))
212 err = tree.Add(k, v)
213 c.Assert(err, qt.IsNil)
214
215 k = BigIntToBytes(big.NewInt(int64(515)))
216 err = tree.Add(k, v)
217 c.Assert(err, qt.IsNil)
218 k = BigIntToBytes(big.NewInt(int64(770)))
219 err = tree.Add(k, v)
220 c.Assert(err, qt.IsNil)
221}
222
223func TestGet(t *testing.T) {
224 c := qt.New(t)
225 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
226 c.Assert(err, qt.IsNil)
227 defer tree.db.Close()
228
229 for i := 0; i < 10; i++ {
230 k := BigIntToBytes(big.NewInt(int64(i)))
231 v := BigIntToBytes(big.NewInt(int64(i * 2)))
232 if err := tree.Add(k, v); err != nil {
233 t.Fatal(err)
234 }
235 }
236
237 k := BigIntToBytes(big.NewInt(int64(7)))
238 gettedKey, gettedValue, err := tree.Get(k)
239 c.Assert(err, qt.IsNil)
240 c.Check(gettedKey, qt.DeepEquals, k)
241 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(big.NewInt(int64(7*2))))
242}
243
244func TestGenProofAndVerify(t *testing.T) {
245 c := qt.New(t)
246 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
247 c.Assert(err, qt.IsNil)
248 defer tree.db.Close()
249
250 for i := 0; i < 10; i++ {
251 k := BigIntToBytes(big.NewInt(int64(i)))
252 v := BigIntToBytes(big.NewInt(int64(i * 2)))
253 if err := tree.Add(k, v); err != nil {
254 t.Fatal(err)
255 }
256 }
257
258 k := BigIntToBytes(big.NewInt(int64(7)))
259 siblings, err := tree.GenProof(k)
260 c.Assert(err, qt.IsNil)
261
262 k = BigIntToBytes(big.NewInt(int64(7)))
263 v := BigIntToBytes(big.NewInt(int64(14)))
264 verif, err := CheckProof(tree.hashFunction, k, v, tree.Root(), siblings)
265 c.Assert(err, qt.IsNil)
266 c.Check(verif, qt.IsTrue)
267}
268
269func TestDumpAndImportDump(t *testing.T) {
270 c := qt.New(t)
271 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
272 c.Assert(err, qt.IsNil)
273 defer tree1.db.Close()
274
275 for i := 0; i < 16; i++ {
276 k := BigIntToBytes(big.NewInt(int64(i)))
277 v := BigIntToBytes(big.NewInt(int64(i * 2)))
278 if err := tree1.Add(k, v); err != nil {
279 t.Fatal(err)
280 }
281 }
282
283 e, err := tree1.Dump()
284 c.Assert(err, qt.IsNil)
285
286 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
287 c.Assert(err, qt.IsNil)
288 defer tree2.db.Close()
289 err = tree2.ImportDump(e)
290 c.Assert(err, qt.IsNil)
291 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
292 c.Check(hex.EncodeToString(tree2.Root()), qt.Equals,
293 "0d93aaa3362b2f999f15e15728f123087c2eee716f01c01f56e23aae07f09f08")
294}
295
296func BenchmarkAdd(b *testing.B) {
297 // prepare inputs
298 var ks, vs [][]byte
299 for i := 0; i < 1000; i++ {
300 k := BigIntToBytes(big.NewInt(int64(i)))
301 v := BigIntToBytes(big.NewInt(int64(i)))
302 ks = append(ks, k)
303 vs = append(vs, v)
304 }
305
306 b.Run("Poseidon", func(b *testing.B) {
307 benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
308 })
309 b.Run("Sha256", func(b *testing.B) {
310 benchmarkAdd(b, HashFunctionSha256, ks, vs)
311 })
312}
313
314func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
315 c := qt.New(b)
316 tree, err := NewTree(memory.NewMemoryStorage(), 140, hashFunc)
317 c.Assert(err, qt.IsNil)
318 defer tree.db.Close()
319
320 for i := 0; i < len(ks); i++ {
321 if err := tree.Add(ks[i], vs[i]); err != nil {
322 b.Fatal(err)
323 }
324 }
325}