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