arbo/tree_test.go
tree_test.goRaw
1package arbo
2
3import (
4 "encoding/hex"
5 "fmt"
6 "math/big"
7 "testing"
8
9 "github.com/iden3/go-merkletree/db/memory"
10 "github.com/stretchr/testify/assert"
11 "github.com/stretchr/testify/require"
12)
13
14func TestAddTestVectors(t *testing.T) {
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(t, HashFunctionPoseidon, testVectorsPoseidon)
23
24 testVectorsSha256 := []string{
25 "0000000000000000000000000000000000000000000000000000000000000000",
26 "46910109172468462938850740851377282682950237270676610513794735904325820156367",
27 "59481735341404520835410489183267411392292882901306595567679529387376287440550",
28 "20573794434149960984975763118181266662429997821552560184909083010514790081771",
29 }
30 testAdd(t, HashFunctionSha256, testVectorsSha256)
31}
32
33func testAdd(t *testing.T, hashFunc HashFunction, testVectors []string) {
34 tree, err := NewTree(memory.NewMemoryStorage(), 10, hashFunc)
35 assert.Nil(t, err)
36 defer tree.db.Close()
37
38 assert.Equal(t, testVectors[0], hex.EncodeToString(tree.Root()))
39
40 err = tree.Add(
41 BigIntToBytes(big.NewInt(1)),
42 BigIntToBytes(big.NewInt(2)))
43 assert.Nil(t, err)
44 rootBI := BytesToBigInt(tree.Root())
45 assert.Equal(t, testVectors[1], rootBI.String())
46
47 err = tree.Add(
48 BigIntToBytes(big.NewInt(33)),
49 BigIntToBytes(big.NewInt(44)))
50 assert.Nil(t, err)
51 rootBI = BytesToBigInt(tree.Root())
52 assert.Equal(t, testVectors[2], rootBI.String())
53
54 err = tree.Add(
55 BigIntToBytes(big.NewInt(1234)),
56 BigIntToBytes(big.NewInt(9876)))
57 assert.Nil(t, err)
58 rootBI = BytesToBigInt(tree.Root())
59 assert.Equal(t, testVectors[3], rootBI.String())
60}
61
62func TestAddBatch(t *testing.T) {
63 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
64 require.Nil(t, err)
65 defer tree.db.Close()
66
67 for i := 0; i < 1000; i++ {
68 k := BigIntToBytes(big.NewInt(int64(i)))
69 v := BigIntToBytes(big.NewInt(0))
70 if err := tree.Add(k, v); err != nil {
71 t.Fatal(err)
72 }
73 }
74
75 rootBI := BytesToBigInt(tree.Root())
76 assert.Equal(t,
77 "296519252211642170490407814696803112091039265640052570497930797516015811235",
78 rootBI.String())
79
80 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
81 require.Nil(t, err)
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 assert.Nil(t, err)
93 assert.Equal(t, 0, len(indexes))
94
95 rootBI = BytesToBigInt(tree2.Root())
96 assert.Equal(t,
97 "296519252211642170490407814696803112091039265640052570497930797516015811235",
98 rootBI.String())
99}
100
101func TestAddDifferentOrder(t *testing.T) {
102 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
103 require.Nil(t, err)
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 require.Nil(t, err)
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 assert.Equal(t, hex.EncodeToString(tree1.Root()), hex.EncodeToString(tree2.Root()))
127 assert.Equal(t,
128 "3b89100bec24da9275c87bc188740389e1d5accfc7d88ba5688d7fa96a00d82f",
129 hex.EncodeToString(tree1.Root()))
130}
131
132func TestAddRepeatedIndex(t *testing.T) {
133 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
134 require.Nil(t, err)
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 assert.NotNil(t, err)
144 assert.Equal(t, fmt.Errorf("max virtual level 100"), err)
145}
146
147func TestAux(t *testing.T) {
148 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
149 require.Nil(t, err)
150 defer tree.db.Close()
151
152 k := BigIntToBytes(big.NewInt(int64(1)))
153 v := BigIntToBytes(big.NewInt(int64(0)))
154 err = tree.Add(k, v)
155 assert.Nil(t, err)
156 k = BigIntToBytes(big.NewInt(int64(256)))
157 err = tree.Add(k, v)
158 assert.Nil(t, err)
159
160 k = BigIntToBytes(big.NewInt(int64(257)))
161 err = tree.Add(k, v)
162 assert.Nil(t, err)
163
164 k = BigIntToBytes(big.NewInt(int64(515)))
165 err = tree.Add(k, v)
166 assert.Nil(t, err)
167 k = BigIntToBytes(big.NewInt(int64(770)))
168 err = tree.Add(k, v)
169 assert.Nil(t, err)
170}
171
172func TestGet(t *testing.T) {
173 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
174 require.Nil(t, err)
175 defer tree.db.Close()
176
177 for i := 0; i < 10; i++ {
178 k := BigIntToBytes(big.NewInt(int64(i)))
179 v := BigIntToBytes(big.NewInt(int64(i * 2)))
180 if err := tree.Add(k, v); err != nil {
181 t.Fatal(err)
182 }
183 }
184
185 k := BigIntToBytes(big.NewInt(int64(7)))
186 gettedKey, gettedValue, err := tree.Get(k)
187 assert.Nil(t, err)
188 assert.Equal(t, k, gettedKey)
189 assert.Equal(t, BigIntToBytes(big.NewInt(int64(7*2))), gettedValue)
190}
191
192func TestGenProofAndVerify(t *testing.T) {
193 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
194 require.Nil(t, err)
195 defer tree.db.Close()
196
197 for i := 0; i < 10; i++ {
198 k := BigIntToBytes(big.NewInt(int64(i)))
199 v := BigIntToBytes(big.NewInt(int64(i * 2)))
200 if err := tree.Add(k, v); err != nil {
201 t.Fatal(err)
202 }
203 }
204
205 k := BigIntToBytes(big.NewInt(int64(7)))
206 siblings, err := tree.GenProof(k)
207 assert.Nil(t, err)
208
209 k = BigIntToBytes(big.NewInt(int64(7)))
210 v := BigIntToBytes(big.NewInt(int64(14)))
211 verif, err := CheckProof(tree.hashFunction, k, v, tree.Root(), siblings)
212 require.Nil(t, err)
213 assert.True(t, verif)
214}
215
216func TestDumpAndImportDump(t *testing.T) {
217 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
218 require.Nil(t, err)
219 defer tree1.db.Close()
220
221 for i := 0; i < 16; i++ {
222 k := BigIntToBytes(big.NewInt(int64(i)))
223 v := BigIntToBytes(big.NewInt(int64(i * 2)))
224 if err := tree1.Add(k, v); err != nil {
225 t.Fatal(err)
226 }
227 }
228
229 e, err := tree1.Dump()
230 require.Nil(t, err)
231
232 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
233 require.Nil(t, err)
234 defer tree2.db.Close()
235 err = tree2.ImportDump(e)
236 require.Nil(t, err)
237 assert.Equal(t, tree1.Root(), tree2.Root())
238 assert.Equal(t,
239 "0d93aaa3362b2f999f15e15728f123087c2eee716f01c01f56e23aae07f09f08",
240 hex.EncodeToString(tree2.Root()))
241}
242
243func BenchmarkAdd(b *testing.B) {
244 // prepare inputs
245 var ks, vs [][]byte
246 for i := 0; i < 1000; i++ {
247 k := BigIntToBytes(big.NewInt(int64(i)))
248 v := BigIntToBytes(big.NewInt(int64(i)))
249 ks = append(ks, k)
250 vs = append(vs, v)
251 }
252
253 b.Run("Poseidon", func(b *testing.B) {
254 benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
255 })
256 b.Run("Sha256", func(b *testing.B) {
257 benchmarkAdd(b, HashFunctionSha256, ks, vs)
258 })
259}
260
261func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
262 tree, err := NewTree(memory.NewMemoryStorage(), 140, hashFunc)
263 require.Nil(b, err)
264 defer tree.db.Close()
265
266 for i := 0; i < len(ks); i++ {
267 if err := tree.Add(ks[i], vs[i]); err != nil {
268 b.Fatal(err)
269 }
270 }
271}