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 assert.Equal(t, testVectors[0], hex.EncodeToString(tree.Root()))
38
39 err = tree.Add(
40 BigIntToBytes(big.NewInt(1)),
41 BigIntToBytes(big.NewInt(2)))
42 assert.Nil(t, err)
43 rootBI := BytesToBigInt(tree.Root())
44 assert.Equal(t, testVectors[1], rootBI.String())
45
46 err = tree.Add(
47 BigIntToBytes(big.NewInt(33)),
48 BigIntToBytes(big.NewInt(44)))
49 assert.Nil(t, err)
50 rootBI = BytesToBigInt(tree.Root())
51 assert.Equal(t, testVectors[2], rootBI.String())
52
53 err = tree.Add(
54 BigIntToBytes(big.NewInt(1234)),
55 BigIntToBytes(big.NewInt(9876)))
56 assert.Nil(t, err)
57 rootBI = BytesToBigInt(tree.Root())
58 assert.Equal(t, testVectors[3], rootBI.String())
59}
60
61func TestAddBatch(t *testing.T) {
62 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
63 require.Nil(t, err)
64 defer tree.db.Close()
65
66 for i := 0; i < 1000; i++ {
67 k := BigIntToBytes(big.NewInt(int64(i)))
68 v := BigIntToBytes(big.NewInt(0))
69 if err := tree.Add(k, v); err != nil {
70 t.Fatal(err)
71 }
72 }
73
74 rootBI := BytesToBigInt(tree.Root())
75 assert.Equal(t,
76 "296519252211642170490407814696803112091039265640052570497930797516015811235",
77 rootBI.String())
78
79 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
80 require.Nil(t, err)
81 defer tree2.db.Close()
82
83 var keys, values [][]byte
84 for i := 0; i < 1000; i++ {
85 k := BigIntToBytes(big.NewInt(int64(i)))
86 v := BigIntToBytes(big.NewInt(0))
87 keys = append(keys, k)
88 values = append(values, v)
89 }
90 indexes, err := tree2.AddBatch(keys, values)
91 assert.Nil(t, err)
92 assert.Equal(t, 0, len(indexes))
93
94 rootBI = BytesToBigInt(tree2.Root())
95 assert.Equal(t,
96 "296519252211642170490407814696803112091039265640052570497930797516015811235",
97 rootBI.String())
98}
99
100func TestAddDifferentOrder(t *testing.T) {
101 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
102 require.Nil(t, err)
103
104 defer tree1.db.Close()
105 for i := 0; i < 16; i++ {
106 k := SwapEndianness(big.NewInt(int64(i)).Bytes())
107 v := SwapEndianness(big.NewInt(0).Bytes())
108 if err := tree1.Add(k, v); err != nil {
109 t.Fatal(err)
110 }
111 }
112
113 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
114 require.Nil(t, err)
115 defer tree2.db.Close()
116 for i := 16 - 1; i >= 0; i-- {
117 k := big.NewInt(int64(i)).Bytes()
118 v := big.NewInt(0).Bytes()
119 if err := tree2.Add(k, v); err != nil {
120 t.Fatal(err)
121 }
122 }
123
124 assert.Equal(t, hex.EncodeToString(tree1.Root()), hex.EncodeToString(tree2.Root()))
125 assert.Equal(t,
126 "3b89100bec24da9275c87bc188740389e1d5accfc7d88ba5688d7fa96a00d82f",
127 hex.EncodeToString(tree1.Root()))
128}
129
130func TestAddRepeatedIndex(t *testing.T) {
131 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
132 require.Nil(t, err)
133 defer tree.db.Close()
134 k := big.NewInt(int64(3)).Bytes()
135 v := big.NewInt(int64(12)).Bytes()
136 if err := tree.Add(k, v); err != nil {
137 t.Fatal(err)
138 }
139 err = tree.Add(k, v)
140 assert.NotNil(t, err)
141 assert.Equal(t, fmt.Errorf("max virtual level 100"), err)
142}
143
144func TestAux(t *testing.T) {
145 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
146 require.Nil(t, err)
147 defer tree.db.Close()
148 k := BigIntToBytes(big.NewInt(int64(1)))
149 v := BigIntToBytes(big.NewInt(int64(0)))
150 err = tree.Add(k, v)
151 assert.Nil(t, err)
152 k = BigIntToBytes(big.NewInt(int64(256)))
153 err = tree.Add(k, v)
154 assert.Nil(t, err)
155
156 k = BigIntToBytes(big.NewInt(int64(257)))
157 err = tree.Add(k, v)
158 assert.Nil(t, err)
159
160 k = BigIntToBytes(big.NewInt(int64(515)))
161 err = tree.Add(k, v)
162 assert.Nil(t, err)
163 k = BigIntToBytes(big.NewInt(int64(770)))
164 err = tree.Add(k, v)
165 assert.Nil(t, err)
166}
167
168func TestGet(t *testing.T) {
169 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
170 require.Nil(t, err)
171
172 defer tree.db.Close()
173 for i := 0; i < 10; i++ {
174 k := BigIntToBytes(big.NewInt(int64(i)))
175 v := BigIntToBytes(big.NewInt(int64(i * 2)))
176 if err := tree.Add(k, v); err != nil {
177 t.Fatal(err)
178 }
179 }
180
181 k := BigIntToBytes(big.NewInt(int64(7)))
182 gettedKey, gettedValue, err := tree.Get(k)
183 assert.Nil(t, err)
184 assert.Equal(t, k, gettedKey)
185 assert.Equal(t, BigIntToBytes(big.NewInt(int64(7*2))), gettedValue)
186}
187
188func TestGenProofAndVerify(t *testing.T) {
189 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
190 require.Nil(t, err)
191
192 defer tree.db.Close()
193 for i := 0; i < 10; i++ {
194 k := BigIntToBytes(big.NewInt(int64(i)))
195 v := BigIntToBytes(big.NewInt(int64(i * 2)))
196 if err := tree.Add(k, v); err != nil {
197 t.Fatal(err)
198 }
199 }
200
201 k := BigIntToBytes(big.NewInt(int64(7)))
202 siblings, err := tree.GenProof(k)
203 assert.Nil(t, err)
204
205 k = BigIntToBytes(big.NewInt(int64(7)))
206 v := BigIntToBytes(big.NewInt(int64(14)))
207 verif, err := CheckProof(tree.hashFunction, k, v, tree.Root(), siblings)
208 require.Nil(t, err)
209 assert.True(t, verif)
210}
211
212func BenchmarkAdd(b *testing.B) {
213 // prepare inputs
214 var ks, vs [][]byte
215 for i := 0; i < 1000; i++ {
216 k := BigIntToBytes(big.NewInt(int64(i)))
217 v := BigIntToBytes(big.NewInt(int64(i)))
218 ks = append(ks, k)
219 vs = append(vs, v)
220 }
221
222 b.Run("Poseidon", func(b *testing.B) {
223 benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
224 })
225 b.Run("Sha256", func(b *testing.B) {
226 benchmarkAdd(b, HashFunctionSha256, ks, vs)
227 })
228}
229
230func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
231 tree, err := NewTree(memory.NewMemoryStorage(), 140, hashFunc)
232 require.Nil(b, err)
233
234 defer tree.db.Close()
235 for i := 0; i < len(ks); i++ {
236 if err := tree.Add(ks[i], vs[i]); err != nil {
237 b.Fatal(err)
238 }
239 }
240}