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