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 "go.vocdoni.io/dvote/db/badgerdb"
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 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
80 c.Assert(err, qt.IsNil)
81 tree, err := NewTree(database, maxLevels, HashFunctionSha256)
82 c.Assert(err, qt.IsNil)
83 for i := 0; i < len(keys); i++ {
84 err := tree.Add(keys[i], values[i])
85 c.Assert(err, qt.IsNil)
86 }
87
88 // virtual tree
89 vTree := newVT(maxLevels, HashFunctionSha256)
90
91 c.Assert(vTree.root, qt.IsNil)
92
93 for i := 0; i < len(keys); i++ {
94 err := vTree.add(0, keys[i], values[i])
95 c.Assert(err, qt.IsNil)
96 }
97
98 // compute hashes, and check Root
99 _, err = vTree.computeHashes()
100 c.Assert(err, qt.IsNil)
101 root, err := tree.Root()
102 c.Assert(err, qt.IsNil)
103 c.Assert(vTree.root.h, qt.DeepEquals, root)
104}
105
106func TestVirtualTreeAddBatch(t *testing.T) {
107 c := qt.New(t)
108
109 nLeafs := 2000
110 maxLevels := 100
111
112 keys := make([][]byte, nLeafs)
113 values := make([][]byte, nLeafs)
114 for i := 0; i < nLeafs; i++ {
115 keys[i] = randomBytes(32)
116 values[i] = randomBytes(32)
117 }
118
119 // normal tree, to have an expected root value
120 database, err := badgerdb.New(badgerdb.Options{Path: c.TempDir()})
121 c.Assert(err, qt.IsNil)
122 tree, err := NewTree(database, maxLevels, HashFunctionBlake2b)
123 c.Assert(err, qt.IsNil)
124 for i := 0; i < len(keys); i++ {
125 err := tree.Add(keys[i], values[i])
126 c.Assert(err, qt.IsNil)
127 }
128
129 // virtual tree
130 vTree := newVT(maxLevels, HashFunctionBlake2b)
131
132 c.Assert(vTree.root, qt.IsNil)
133
134 invalids, err := vTree.addBatch(keys, values)
135 c.Assert(err, qt.IsNil)
136 c.Assert(len(invalids), qt.Equals, 0)
137
138 // compute hashes, and check Root
139 _, err = vTree.computeHashes()
140 c.Assert(err, qt.IsNil)
141 root, err := tree.Root()
142 c.Assert(err, qt.IsNil)
143 c.Assert(vTree.root.h, qt.DeepEquals, root)
144}
145
146func TestVirtualTreeAddBatchFullyUsed(t *testing.T) {
147 c := qt.New(t)
148
149 vTree1 := newVT(7, HashFunctionPoseidon) // used for add one by one
150 vTree2 := newVT(7, HashFunctionPoseidon) // used for addBatch
151
152 var keys, values [][]byte
153 for i := 0; i < 128; i++ {
154 k := BigIntToBytes(32, big.NewInt(int64(i)))
155 v := k
156
157 keys = append(keys, k)
158 values = append(values, v)
159
160 // add one by one expecting no error
161 err := vTree1.add(0, k, v)
162 c.Assert(err, qt.IsNil)
163 }
164
165 invalids, err := vTree2.addBatch(keys, values)
166 c.Assert(err, qt.IsNil)
167 c.Assert(0, qt.Equals, len(invalids))
168}
169
170func TestGetNodesAtLevel(t *testing.T) {
171 c := qt.New(t)
172
173 tree0 := vt{
174 params: ¶ms{
175 maxLevels: 100,
176 hashFunction: HashFunctionBlake2b,
177 emptyHash: make([]byte, HashFunctionBlake2b.Len()),
178 },
179 root: nil,
180 }
181
182 tree1 := vt{
183 params: ¶ms{
184 maxLevels: 100,
185 hashFunction: HashFunctionBlake2b,
186 emptyHash: make([]byte, HashFunctionBlake2b.Len()),
187 },
188 root: &node{
189 l: &node{
190 l: &node{
191 k: []byte{0, 0, 0, 0},
192 v: []byte{0, 0, 0, 0},
193 },
194 r: &node{
195 k: []byte{0, 0, 0, 1},
196 v: []byte{0, 0, 0, 1},
197 },
198 },
199 r: &node{
200 l: &node{
201 k: []byte{0, 0, 0, 2},
202 v: []byte{0, 0, 0, 2},
203 },
204 r: &node{
205 k: []byte{0, 0, 0, 3},
206 v: []byte{0, 0, 0, 3},
207 },
208 },
209 },
210 }
211 // tree1.printGraphviz()
212
213 tree2 := vt{
214 params: ¶ms{
215 maxLevels: 100,
216 hashFunction: HashFunctionBlake2b,
217 emptyHash: make([]byte, HashFunctionBlake2b.Len()),
218 },
219 root: &node{
220 l: nil,
221 r: &node{
222 l: &node{
223 l: &node{
224 l: &node{
225 k: []byte{0, 0, 0, 0},
226 v: []byte{0, 0, 0, 0},
227 },
228 r: &node{
229 k: []byte{0, 0, 0, 1},
230 v: []byte{0, 0, 0, 1},
231 },
232 },
233 r: &node{
234 k: []byte{0, 0, 0, 2},
235 v: []byte{0, 0, 0, 2},
236 },
237 },
238 r: &node{
239 k: []byte{0, 0, 0, 3},
240 v: []byte{0, 0, 0, 3},
241 },
242 },
243 },
244 }
245 // tree2.printGraphviz()
246
247 tree3 := vt{
248 params: ¶ms{
249 maxLevels: 100,
250 hashFunction: HashFunctionBlake2b,
251 emptyHash: make([]byte, HashFunctionBlake2b.Len()),
252 },
253 root: &node{
254 l: nil,
255 r: &node{
256 l: &node{
257 l: &node{
258 l: &node{
259 k: []byte{0, 0, 0, 0},
260 v: []byte{0, 0, 0, 0},
261 },
262 r: &node{
263 k: []byte{0, 0, 0, 1},
264 v: []byte{0, 0, 0, 1},
265 },
266 },
267 r: &node{
268 k: []byte{0, 0, 0, 2},
269 v: []byte{0, 0, 0, 2},
270 },
271 },
272 r: nil,
273 },
274 },
275 }
276 // tree3.printGraphviz()
277
278 nodes0, err := tree0.getNodesAtLevel(2)
279 c.Assert(err, qt.IsNil)
280 c.Assert(len(nodes0), qt.DeepEquals, 4)
281 c.Assert("0000", qt.DeepEquals, getNotNils(nodes0))
282
283 nodes1, err := tree1.getNodesAtLevel(2)
284 c.Assert(err, qt.IsNil)
285 c.Assert(len(nodes1), qt.DeepEquals, 4)
286 c.Assert("1111", qt.DeepEquals, getNotNils(nodes1))
287
288 nodes1, err = tree1.getNodesAtLevel(3)
289 c.Assert(err, qt.IsNil)
290 c.Assert(len(nodes1), qt.DeepEquals, 8)
291 c.Assert("00000000", qt.DeepEquals, getNotNils(nodes1))
292
293 nodes2, err := tree2.getNodesAtLevel(2)
294 c.Assert(err, qt.IsNil)
295 c.Assert(len(nodes2), qt.DeepEquals, 4)
296 c.Assert("0011", qt.DeepEquals, getNotNils(nodes2))
297
298 nodes2, err = tree2.getNodesAtLevel(3)
299 c.Assert(err, qt.IsNil)
300 c.Assert(len(nodes2), qt.DeepEquals, 8)
301 c.Assert("00001100", qt.DeepEquals, getNotNils(nodes2))
302
303 nodes3, err := tree3.getNodesAtLevel(2)
304 c.Assert(err, qt.IsNil)
305 c.Assert(len(nodes3), qt.DeepEquals, 4)
306 c.Assert("0010", qt.DeepEquals, getNotNils(nodes3))
307
308 nodes3, err = tree3.getNodesAtLevel(3)
309 c.Assert(err, qt.IsNil)
310 c.Assert(len(nodes3), qt.DeepEquals, 8)
311 c.Assert("00001100", qt.DeepEquals, getNotNils(nodes3))
312
313 nodes3, err = tree3.getNodesAtLevel(4)
314 c.Assert(err, qt.IsNil)
315 c.Assert(len(nodes3), qt.DeepEquals, 16)
316 c.Assert("0000000011000000", qt.DeepEquals, getNotNils(nodes3))
317}
318
319func getNotNils(nodes []*node) string {
320 s := ""
321 for i := 0; i < len(nodes); i++ {
322 if nodes[i] == nil {
323 s += "0"
324 } else {
325 s += "1"
326 }
327 }
328 return s
329}