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