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 := SwapEndianness(big.NewInt(int64(i)).Bytes())
108 v := SwapEndianness(big.NewInt(0).Bytes())
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 := big.NewInt(int64(i)).Bytes()
120 v := big.NewInt(0).Bytes()
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 := big.NewInt(int64(3)).Bytes()
138 v := big.NewInt(int64(12)).Bytes()
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 TestAux(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(1)))
154 v := BigIntToBytes(big.NewInt(int64(0)))
155 err = tree.Add(k, v)
156 c.Assert(err, qt.IsNil)
157 k = BigIntToBytes(big.NewInt(int64(256)))
158 err = tree.Add(k, v)
159 c.Assert(err, qt.IsNil)
160
161 k = BigIntToBytes(big.NewInt(int64(257)))
162 err = tree.Add(k, v)
163 c.Assert(err, qt.IsNil)
164
165 k = BigIntToBytes(big.NewInt(int64(515)))
166 err = tree.Add(k, v)
167 c.Assert(err, qt.IsNil)
168 k = BigIntToBytes(big.NewInt(int64(770)))
169 err = tree.Add(k, v)
170 c.Assert(err, qt.IsNil)
171}
172
173func TestGet(t *testing.T) {
174 c := qt.New(t)
175 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
176 c.Assert(err, qt.IsNil)
177 defer tree.db.Close()
178
179 for i := 0; i < 10; i++ {
180 k := BigIntToBytes(big.NewInt(int64(i)))
181 v := BigIntToBytes(big.NewInt(int64(i * 2)))
182 if err := tree.Add(k, v); err != nil {
183 t.Fatal(err)
184 }
185 }
186
187 k := BigIntToBytes(big.NewInt(int64(7)))
188 gettedKey, gettedValue, err := tree.Get(k)
189 c.Assert(err, qt.IsNil)
190 c.Check(gettedKey, qt.DeepEquals, k)
191 c.Check(gettedValue, qt.DeepEquals, BigIntToBytes(big.NewInt(int64(7*2))))
192}
193
194func TestGenProofAndVerify(t *testing.T) {
195 c := qt.New(t)
196 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
197 c.Assert(err, qt.IsNil)
198 defer tree.db.Close()
199
200 for i := 0; i < 10; i++ {
201 k := BigIntToBytes(big.NewInt(int64(i)))
202 v := BigIntToBytes(big.NewInt(int64(i * 2)))
203 if err := tree.Add(k, v); err != nil {
204 t.Fatal(err)
205 }
206 }
207
208 k := BigIntToBytes(big.NewInt(int64(7)))
209 siblings, err := tree.GenProof(k)
210 c.Assert(err, qt.IsNil)
211
212 k = BigIntToBytes(big.NewInt(int64(7)))
213 v := BigIntToBytes(big.NewInt(int64(14)))
214 verif, err := CheckProof(tree.hashFunction, k, v, tree.Root(), siblings)
215 c.Assert(err, qt.IsNil)
216 c.Check(verif, qt.IsTrue)
217}
218
219func TestDumpAndImportDump(t *testing.T) {
220 c := qt.New(t)
221 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
222 c.Assert(err, qt.IsNil)
223 defer tree1.db.Close()
224
225 for i := 0; i < 16; i++ {
226 k := BigIntToBytes(big.NewInt(int64(i)))
227 v := BigIntToBytes(big.NewInt(int64(i * 2)))
228 if err := tree1.Add(k, v); err != nil {
229 t.Fatal(err)
230 }
231 }
232
233 e, err := tree1.Dump()
234 c.Assert(err, qt.IsNil)
235
236 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
237 c.Assert(err, qt.IsNil)
238 defer tree2.db.Close()
239 err = tree2.ImportDump(e)
240 c.Assert(err, qt.IsNil)
241 c.Check(tree2.Root(), qt.DeepEquals, tree1.Root())
242 c.Check(hex.EncodeToString(tree2.Root()), qt.Equals,
243 "0d93aaa3362b2f999f15e15728f123087c2eee716f01c01f56e23aae07f09f08")
244}
245
246func BenchmarkAdd(b *testing.B) {
247 // prepare inputs
248 var ks, vs [][]byte
249 for i := 0; i < 1000; i++ {
250 k := BigIntToBytes(big.NewInt(int64(i)))
251 v := BigIntToBytes(big.NewInt(int64(i)))
252 ks = append(ks, k)
253 vs = append(vs, v)
254 }
255
256 b.Run("Poseidon", func(b *testing.B) {
257 benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
258 })
259 b.Run("Sha256", func(b *testing.B) {
260 benchmarkAdd(b, HashFunctionSha256, ks, vs)
261 })
262}
263
264func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
265 c := qt.New(b)
266 tree, err := NewTree(memory.NewMemoryStorage(), 140, hashFunc)
267 c.Assert(err, qt.IsNil)
268 defer tree.db.Close()
269
270 for i := 0; i < len(ks); i++ {
271 if err := tree.Add(ks[i], vs[i]); err != nil {
272 b.Fatal(err)
273 }
274 }
275}