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