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