arbo/vt_test.go
vt_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 TestVirtualTreeTestVectors(t *testing.T) {
13 c := qt.New(t)
14
15 bLen := 32
16 keys := [][]byte{
17 BigIntToBytes(bLen, big.NewInt(1)),
18 BigIntToBytes(bLen, big.NewInt(33)),
19 BigIntToBytes(bLen, big.NewInt(1234)),
20 BigIntToBytes(bLen, big.NewInt(123456789)),
21 }
22 values := [][]byte{
23 BigIntToBytes(bLen, big.NewInt(2)),
24 BigIntToBytes(bLen, big.NewInt(44)),
25 BigIntToBytes(bLen, big.NewInt(9876)),
26 BigIntToBytes(bLen, big.NewInt(987654321)),
27 }
28
29 // check the root for different batches of leafs
30 testVirtualTree(c, 10, keys[:1], values[:1])
31 testVirtualTree(c, 10, keys[:2], values[:2])
32 testVirtualTree(c, 10, keys[:3], values[:3])
33 testVirtualTree(c, 10, keys[:4], values[:4])
34
35 // test with hardcoded values
36 testvectorKeys := []string{
37 "1c7c2265e368314ca58ed2e1f33a326f1220e234a566d55c3605439dbe411642",
38 "2c9f0a578afff5bfa4e0992a43066460faaab9e8e500db0b16647c701cdb16bf",
39 "9cb87ec67e875c61390edcd1ab517f443591047709a4d4e45b0f9ed980857b8e",
40 "9b4e9e92e974a589f426ceeb4cb291dc24893513fecf8e8460992dcf52621d4d",
41 "1c45cb31f2fa39ec7b9ebf0fad40e0b8296016b5ce8844ae06ff77226379d9a5",
42 "d8af98bbbb585129798ae54d5eabbc9d0561d583faf1663b3a3724d15bda4ec7",
43 "3cd55dbfb8f975f20a0925dfbdabe79fa2d51dd0268afbb8ba6b01de9dfcdd3c",
44 "5d0a9d6d9f197c091bf054fac9cb60e11ec723d6610ed8578e617b4d46cb43d5",
45 }
46 keys = [][]byte{}
47 values = [][]byte{}
48 for i := 0; i < len(testvectorKeys); i++ {
49 key, err := hex.DecodeString(testvectorKeys[i])
50 c.Assert(err, qt.IsNil)
51 keys = append(keys, key)
52 values = append(values, []byte{0})
53 }
54
55 // check the root for different batches of leafs
56 testVirtualTree(c, 10, keys[:1], values[:1])
57 testVirtualTree(c, 10, keys, values)
58}
59
60func TestVirtualTreeRandomKeys(t *testing.T) {
61 c := qt.New(t)
62
63 // test with random values
64 nLeafs := 1024
65 keys := make([][]byte, nLeafs)
66 values := make([][]byte, nLeafs)
67 for i := 0; i < nLeafs; i++ {
68 keys[i] = randomBytes(32)
69 values[i] = randomBytes(32)
70 }
71
72 testVirtualTree(c, 100, keys, values)
73}
74
75func testVirtualTree(c *qt.C, maxLevels int, keys, values [][]byte) {
76 c.Assert(len(keys), qt.Equals, len(values))
77
78 // normal tree, to have an expected root value
79 tree, err := NewTree(memory.NewMemoryStorage(), maxLevels, HashFunctionSha256)
80 c.Assert(err, qt.IsNil)
81 for i := 0; i < len(keys); i++ {
82 err := tree.Add(keys[i], values[i])
83 c.Assert(err, qt.IsNil)
84 }
85
86 // virtual tree
87 vTree := newVT(maxLevels, HashFunctionSha256)
88
89 c.Assert(vTree.root, qt.IsNil)
90
91 for i := 0; i < len(keys); i++ {
92 err := vTree.add(0, keys[i], values[i])
93 c.Assert(err, qt.IsNil)
94 }
95
96 // compute hashes, and check Root
97 _, err = vTree.computeHashes()
98 c.Assert(err, qt.IsNil)
99 c.Assert(vTree.root.h, qt.DeepEquals, tree.root)
100}
101
102func TestVirtualTreeAddBatch(t *testing.T) {
103 c := qt.New(t)
104
105 nLeafs := 2000
106 maxLevels := 100
107
108 keys := make([][]byte, nLeafs)
109 values := make([][]byte, nLeafs)
110 for i := 0; i < nLeafs; i++ {
111 keys[i] = randomBytes(32)
112 values[i] = randomBytes(32)
113 }
114
115 // normal tree, to have an expected root value
116 tree, err := NewTree(memory.NewMemoryStorage(), maxLevels, HashFunctionBlake2b)
117 c.Assert(err, qt.IsNil)
118 for i := 0; i < len(keys); i++ {
119 err := tree.Add(keys[i], values[i])
120 c.Assert(err, qt.IsNil)
121 }
122
123 // virtual tree
124 vTree := newVT(maxLevels, HashFunctionBlake2b)
125
126 c.Assert(vTree.root, qt.IsNil)
127
128 invalids, err := vTree.addBatch(keys, values)
129 c.Assert(err, qt.IsNil)
130 c.Assert(len(invalids), qt.Equals, 0)
131
132 // compute hashes, and check Root
133 _, err = vTree.computeHashes()
134 c.Assert(err, qt.IsNil)
135 c.Assert(vTree.root.h, qt.DeepEquals, tree.root)
136}
137
138func TestGetNodesAtLevel(t *testing.T) {
139 c := qt.New(t)
140
141 tree0 := vt{
142 params: ¶ms{
143 maxLevels: 100,
144 hashFunction: HashFunctionBlake2b,
145 emptyHash: make([]byte, HashFunctionBlake2b.Len()),
146 },
147 root: nil,
148 }
149
150 tree1 := vt{
151 params: ¶ms{
152 maxLevels: 100,
153 hashFunction: HashFunctionBlake2b,
154 emptyHash: make([]byte, HashFunctionBlake2b.Len()),
155 },
156 root: &node{
157 l: &node{
158 l: &node{
159 k: []byte{0, 0, 0, 0},
160 v: []byte{0, 0, 0, 0},
161 },
162 r: &node{
163 k: []byte{0, 0, 0, 1},
164 v: []byte{0, 0, 0, 1},
165 },
166 },
167 r: &node{
168 l: &node{
169 k: []byte{0, 0, 0, 2},
170 v: []byte{0, 0, 0, 2},
171 },
172 r: &node{
173 k: []byte{0, 0, 0, 3},
174 v: []byte{0, 0, 0, 3},
175 },
176 },
177 },
178 }
179 // tree1.printGraphviz()
180
181 tree2 := vt{
182 params: ¶ms{
183 maxLevels: 100,
184 hashFunction: HashFunctionBlake2b,
185 emptyHash: make([]byte, HashFunctionBlake2b.Len()),
186 },
187 root: &node{
188 l: nil,
189 r: &node{
190 l: &node{
191 l: &node{
192 l: &node{
193 k: []byte{0, 0, 0, 0},
194 v: []byte{0, 0, 0, 0},
195 },
196 r: &node{
197 k: []byte{0, 0, 0, 1},
198 v: []byte{0, 0, 0, 1},
199 },
200 },
201 r: &node{
202 k: []byte{0, 0, 0, 2},
203 v: []byte{0, 0, 0, 2},
204 },
205 },
206 r: &node{
207 k: []byte{0, 0, 0, 3},
208 v: []byte{0, 0, 0, 3},
209 },
210 },
211 },
212 }
213 // tree2.printGraphviz()
214
215 tree3 := vt{
216 params: ¶ms{
217 maxLevels: 100,
218 hashFunction: HashFunctionBlake2b,
219 emptyHash: make([]byte, HashFunctionBlake2b.Len()),
220 },
221 root: &node{
222 l: nil,
223 r: &node{
224 l: &node{
225 l: &node{
226 l: &node{
227 k: []byte{0, 0, 0, 0},
228 v: []byte{0, 0, 0, 0},
229 },
230 r: &node{
231 k: []byte{0, 0, 0, 1},
232 v: []byte{0, 0, 0, 1},
233 },
234 },
235 r: &node{
236 k: []byte{0, 0, 0, 2},
237 v: []byte{0, 0, 0, 2},
238 },
239 },
240 r: nil,
241 },
242 },
243 }
244 // tree3.printGraphviz()
245
246 nodes0, err := tree0.getNodesAtLevel(2)
247 c.Assert(err, qt.IsNil)
248 c.Assert(len(nodes0), qt.DeepEquals, 4)
249 c.Assert("0000", qt.DeepEquals, getNotNils(nodes0))
250
251 nodes1, err := tree1.getNodesAtLevel(2)
252 c.Assert(err, qt.IsNil)
253 c.Assert(len(nodes1), qt.DeepEquals, 4)
254 c.Assert("1111", qt.DeepEquals, getNotNils(nodes1))
255
256 nodes1, err = tree1.getNodesAtLevel(3)
257 c.Assert(err, qt.IsNil)
258 c.Assert(len(nodes1), qt.DeepEquals, 8)
259 c.Assert("00000000", qt.DeepEquals, getNotNils(nodes1))
260
261 nodes2, err := tree2.getNodesAtLevel(2)
262 c.Assert(err, qt.IsNil)
263 c.Assert(len(nodes2), qt.DeepEquals, 4)
264 c.Assert("0011", qt.DeepEquals, getNotNils(nodes2))
265
266 nodes2, err = tree2.getNodesAtLevel(3)
267 c.Assert(err, qt.IsNil)
268 c.Assert(len(nodes2), qt.DeepEquals, 8)
269 c.Assert("00001100", qt.DeepEquals, getNotNils(nodes2))
270
271 nodes3, err := tree3.getNodesAtLevel(2)
272 c.Assert(err, qt.IsNil)
273 c.Assert(len(nodes3), qt.DeepEquals, 4)
274 c.Assert("0010", qt.DeepEquals, getNotNils(nodes3))
275
276 nodes3, err = tree3.getNodesAtLevel(3)
277 c.Assert(err, qt.IsNil)
278 c.Assert(len(nodes3), qt.DeepEquals, 8)
279 c.Assert("00001100", qt.DeepEquals, getNotNils(nodes3))
280
281 nodes3, err = tree3.getNodesAtLevel(4)
282 c.Assert(err, qt.IsNil)
283 c.Assert(len(nodes3), qt.DeepEquals, 16)
284 c.Assert("0000000011000000", qt.DeepEquals, getNotNils(nodes3))
285}
286
287func getNotNils(nodes []*node) string {
288 s := ""
289 for i := 0; i < len(nodes); i++ {
290 if nodes[i] == nil {
291 s += "0"
292 } else {
293 s += "1"
294 }
295 }
296 return s
297}