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 TestGetNodesAtLevel(t *testing.T) {
147 c := qt.New(t)
148
149 tree0 := vt{
150 params: ¶ms{
151 maxLevels: 100,
152 hashFunction: HashFunctionBlake2b,
153 emptyHash: make([]byte, HashFunctionBlake2b.Len()),
154 },
155 root: nil,
156 }
157
158 tree1 := vt{
159 params: ¶ms{
160 maxLevels: 100,
161 hashFunction: HashFunctionBlake2b,
162 emptyHash: make([]byte, HashFunctionBlake2b.Len()),
163 },
164 root: &node{
165 l: &node{
166 l: &node{
167 k: []byte{0, 0, 0, 0},
168 v: []byte{0, 0, 0, 0},
169 },
170 r: &node{
171 k: []byte{0, 0, 0, 1},
172 v: []byte{0, 0, 0, 1},
173 },
174 },
175 r: &node{
176 l: &node{
177 k: []byte{0, 0, 0, 2},
178 v: []byte{0, 0, 0, 2},
179 },
180 r: &node{
181 k: []byte{0, 0, 0, 3},
182 v: []byte{0, 0, 0, 3},
183 },
184 },
185 },
186 }
187 // tree1.printGraphviz()
188
189 tree2 := vt{
190 params: ¶ms{
191 maxLevels: 100,
192 hashFunction: HashFunctionBlake2b,
193 emptyHash: make([]byte, HashFunctionBlake2b.Len()),
194 },
195 root: &node{
196 l: nil,
197 r: &node{
198 l: &node{
199 l: &node{
200 l: &node{
201 k: []byte{0, 0, 0, 0},
202 v: []byte{0, 0, 0, 0},
203 },
204 r: &node{
205 k: []byte{0, 0, 0, 1},
206 v: []byte{0, 0, 0, 1},
207 },
208 },
209 r: &node{
210 k: []byte{0, 0, 0, 2},
211 v: []byte{0, 0, 0, 2},
212 },
213 },
214 r: &node{
215 k: []byte{0, 0, 0, 3},
216 v: []byte{0, 0, 0, 3},
217 },
218 },
219 },
220 }
221 // tree2.printGraphviz()
222
223 tree3 := vt{
224 params: ¶ms{
225 maxLevels: 100,
226 hashFunction: HashFunctionBlake2b,
227 emptyHash: make([]byte, HashFunctionBlake2b.Len()),
228 },
229 root: &node{
230 l: nil,
231 r: &node{
232 l: &node{
233 l: &node{
234 l: &node{
235 k: []byte{0, 0, 0, 0},
236 v: []byte{0, 0, 0, 0},
237 },
238 r: &node{
239 k: []byte{0, 0, 0, 1},
240 v: []byte{0, 0, 0, 1},
241 },
242 },
243 r: &node{
244 k: []byte{0, 0, 0, 2},
245 v: []byte{0, 0, 0, 2},
246 },
247 },
248 r: nil,
249 },
250 },
251 }
252 // tree3.printGraphviz()
253
254 nodes0, err := tree0.getNodesAtLevel(2)
255 c.Assert(err, qt.IsNil)
256 c.Assert(len(nodes0), qt.DeepEquals, 4)
257 c.Assert("0000", qt.DeepEquals, getNotNils(nodes0))
258
259 nodes1, err := tree1.getNodesAtLevel(2)
260 c.Assert(err, qt.IsNil)
261 c.Assert(len(nodes1), qt.DeepEquals, 4)
262 c.Assert("1111", qt.DeepEquals, getNotNils(nodes1))
263
264 nodes1, err = tree1.getNodesAtLevel(3)
265 c.Assert(err, qt.IsNil)
266 c.Assert(len(nodes1), qt.DeepEquals, 8)
267 c.Assert("00000000", qt.DeepEquals, getNotNils(nodes1))
268
269 nodes2, err := tree2.getNodesAtLevel(2)
270 c.Assert(err, qt.IsNil)
271 c.Assert(len(nodes2), qt.DeepEquals, 4)
272 c.Assert("0011", qt.DeepEquals, getNotNils(nodes2))
273
274 nodes2, err = tree2.getNodesAtLevel(3)
275 c.Assert(err, qt.IsNil)
276 c.Assert(len(nodes2), qt.DeepEquals, 8)
277 c.Assert("00001100", qt.DeepEquals, getNotNils(nodes2))
278
279 nodes3, err := tree3.getNodesAtLevel(2)
280 c.Assert(err, qt.IsNil)
281 c.Assert(len(nodes3), qt.DeepEquals, 4)
282 c.Assert("0010", qt.DeepEquals, getNotNils(nodes3))
283
284 nodes3, err = tree3.getNodesAtLevel(3)
285 c.Assert(err, qt.IsNil)
286 c.Assert(len(nodes3), qt.DeepEquals, 8)
287 c.Assert("00001100", qt.DeepEquals, getNotNils(nodes3))
288
289 nodes3, err = tree3.getNodesAtLevel(4)
290 c.Assert(err, qt.IsNil)
291 c.Assert(len(nodes3), qt.DeepEquals, 16)
292 c.Assert("0000000011000000", qt.DeepEquals, getNotNils(nodes3))
293}
294
295func getNotNils(nodes []*node) string {
296 s := ""
297 for i := 0; i < len(nodes); i++ {
298 if nodes[i] == nil {
299 s += "0"
300 } else {
301 s += "1"
302 }
303 }
304 return s
305}