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 TestAdd1000(t *testing.T) {
62 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
63 require.Nil(t, err)
64
65 defer tree.db.Close()
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
80func TestAddDifferentOrder(t *testing.T) {
81 tree1, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
82 require.Nil(t, err)
83
84 defer tree1.db.Close()
85 for i := 0; i < 16; i++ {
86 k := SwapEndianness(big.NewInt(int64(i)).Bytes())
87 v := SwapEndianness(big.NewInt(0).Bytes())
88 if err := tree1.Add(k, v); err != nil {
89 t.Fatal(err)
90 }
91 }
92
93 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
94 require.Nil(t, err)
95 defer tree2.db.Close()
96 for i := 16 - 1; i >= 0; i-- {
97 k := big.NewInt(int64(i)).Bytes()
98 v := big.NewInt(0).Bytes()
99 if err := tree2.Add(k, v); err != nil {
100 t.Fatal(err)
101 }
102 }
103
104 assert.Equal(t, hex.EncodeToString(tree1.Root()), hex.EncodeToString(tree2.Root()))
105 assert.Equal(t,
106 "3b89100bec24da9275c87bc188740389e1d5accfc7d88ba5688d7fa96a00d82f",
107 hex.EncodeToString(tree1.Root()))
108}
109
110func TestAddRepeatedIndex(t *testing.T) {
111 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
112 require.Nil(t, err)
113 defer tree.db.Close()
114 k := big.NewInt(int64(3)).Bytes()
115 v := big.NewInt(int64(12)).Bytes()
116 if err := tree.Add(k, v); err != nil {
117 t.Fatal(err)
118 }
119 err = tree.Add(k, v)
120 assert.NotNil(t, err)
121 assert.Equal(t, fmt.Errorf("max virtual level 100"), err)
122}
123
124func TestAux(t *testing.T) {
125 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
126 require.Nil(t, err)
127 defer tree.db.Close()
128 k := BigIntToBytes(big.NewInt(int64(1)))
129 v := BigIntToBytes(big.NewInt(int64(0)))
130 err = tree.Add(k, v)
131 assert.Nil(t, err)
132 k = BigIntToBytes(big.NewInt(int64(256)))
133 err = tree.Add(k, v)
134 assert.Nil(t, err)
135
136 k = BigIntToBytes(big.NewInt(int64(257)))
137 err = tree.Add(k, v)
138 assert.Nil(t, err)
139
140 k = BigIntToBytes(big.NewInt(int64(515)))
141 err = tree.Add(k, v)
142 assert.Nil(t, err)
143 k = BigIntToBytes(big.NewInt(int64(770)))
144 err = tree.Add(k, v)
145 assert.Nil(t, err)
146}
147
148func TestGenProofAndVerify(t *testing.T) {
149 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
150 require.Nil(t, err)
151
152 defer tree.db.Close()
153 for i := 0; i < 10; i++ {
154 k := BigIntToBytes(big.NewInt(int64(i)))
155 v := BigIntToBytes(big.NewInt(int64(i * 2)))
156 if err := tree.Add(k, v); err != nil {
157 t.Fatal(err)
158 }
159 }
160
161 k := BigIntToBytes(big.NewInt(int64(7)))
162 siblings, err := tree.GenProof(k)
163 assert.Nil(t, err)
164
165 k = BigIntToBytes(big.NewInt(int64(7)))
166 v := BigIntToBytes(big.NewInt(int64(14)))
167 verif, err := CheckProof(tree.hashFunction, k, v, tree.Root(), siblings)
168 require.Nil(t, err)
169 assert.True(t, verif)
170}
171
172func BenchmarkAdd(b *testing.B) {
173 // prepare inputs
174 var ks, vs [][]byte
175 for i := 0; i < 1000; i++ {
176 k := BigIntToBytes(big.NewInt(int64(i)))
177 v := BigIntToBytes(big.NewInt(int64(i)))
178 ks = append(ks, k)
179 vs = append(vs, v)
180 }
181
182 b.Run("Poseidon", func(b *testing.B) {
183 benchmarkAdd(b, HashFunctionPoseidon, ks, vs)
184 })
185 b.Run("Sha256", func(b *testing.B) {
186 benchmarkAdd(b, HashFunctionSha256, ks, vs)
187 })
188}
189
190func benchmarkAdd(b *testing.B, hashFunc HashFunction, ks, vs [][]byte) {
191 tree, err := NewTree(memory.NewMemoryStorage(), 140, hashFunc)
192 require.Nil(b, err)
193
194 defer tree.db.Close()
195 for i := 0; i < len(ks); i++ {
196 if err := tree.Add(ks[i], vs[i]); err != nil {
197 b.Fatal(err)
198 }
199 }
200}