arbo/addbatch_test.go
addbatch_test.goRaw
1package arbo
2
3import (
4 "encoding/hex"
5 "fmt"
6 "math/big"
7 "testing"
8 "time"
9
10 qt "github.com/frankban/quicktest"
11 "github.com/iden3/go-merkletree/db/memory"
12)
13
14func TestAddBatchCaseA(t *testing.T) {
15 c := qt.New(t)
16
17 nLeafs := 1024
18
19 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
20 c.Assert(err, qt.IsNil)
21 defer tree.db.Close()
22
23 start := time.Now()
24 for i := 0; i < nLeafs; i++ {
25 k := BigIntToBytes(big.NewInt(int64(i)))
26 v := BigIntToBytes(big.NewInt(int64(i * 2)))
27 if err := tree.Add(k, v); err != nil {
28 t.Fatal(err)
29 }
30 }
31 fmt.Println(time.Since(start))
32
33 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
34 c.Assert(err, qt.IsNil)
35 defer tree2.db.Close()
36
37 var keys, values [][]byte
38 for i := 0; i < nLeafs; i++ {
39 k := BigIntToBytes(big.NewInt(int64(i)))
40 v := BigIntToBytes(big.NewInt(int64(i * 2)))
41 keys = append(keys, k)
42 values = append(values, v)
43 }
44 start = time.Now()
45 indexes, err := tree2.AddBatchOpt(keys, values)
46 c.Assert(err, qt.IsNil)
47 fmt.Println(time.Since(start))
48 c.Check(len(indexes), qt.Equals, 0)
49
50 // check that both trees roots are equal
51 c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
52}
53
54func TestAddBatchCaseB(t *testing.T) {
55 c := qt.New(t)
56
57 nLeafs := 1024
58
59 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
60 c.Assert(err, qt.IsNil)
61 defer tree.db.Close()
62
63 start := time.Now()
64 for i := 0; i < nLeafs; i++ {
65 k := BigIntToBytes(big.NewInt(int64(i)))
66 v := BigIntToBytes(big.NewInt(int64(i * 2)))
67 if err := tree.Add(k, v); err != nil {
68 t.Fatal(err)
69 }
70 }
71 fmt.Println(time.Since(start))
72
73 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
74 c.Assert(err, qt.IsNil)
75 defer tree2.db.Close()
76
77 // add the initial leafs to fill a bit the tree before calling the
78 // AddBatch method
79 for i := 0; i < 99; i++ { // TMP TODO use const minLeafsThreshold-1 once ready
80 k := BigIntToBytes(big.NewInt(int64(i)))
81 v := BigIntToBytes(big.NewInt(int64(i * 2)))
82 if err := tree2.Add(k, v); err != nil {
83 t.Fatal(err)
84 }
85 }
86
87 var keys, values [][]byte
88 for i := 99; i < nLeafs; i++ {
89 k := BigIntToBytes(big.NewInt(int64(i)))
90 v := BigIntToBytes(big.NewInt(int64(i * 2)))
91 keys = append(keys, k)
92 values = append(values, v)
93 }
94 start = time.Now()
95 indexes, err := tree2.AddBatchOpt(keys, values)
96 c.Assert(err, qt.IsNil)
97 fmt.Println(time.Since(start))
98 c.Check(len(indexes), qt.Equals, 0)
99
100 // check that both trees roots are equal
101 c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
102}
103
104func TestGetKeysAtLevel(t *testing.T) {
105 c := qt.New(t)
106
107 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
108 c.Assert(err, qt.IsNil)
109 defer tree.db.Close()
110
111 for i := 0; i < 32; i++ {
112 k := BigIntToBytes(big.NewInt(int64(i)))
113 v := BigIntToBytes(big.NewInt(int64(i * 2)))
114 if err := tree.Add(k, v); err != nil {
115 t.Fatal(err)
116 }
117 }
118
119 keys, err := tree.getKeysAtLevel(2)
120 c.Assert(err, qt.IsNil)
121 expected := []string{
122 "a5d5f14fce7348e40751496cf25d107d91b0bd043435b9577d778a01f8aa6111",
123 "e9e8dd9b28a7f81d1ff34cb5cefc0146dd848b31031a427b79bdadb62e7f6910",
124 }
125 for i := 0; i < len(keys); i++ {
126 c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
127 }
128
129 keys, err = tree.getKeysAtLevel(3)
130 c.Assert(err, qt.IsNil)
131 expected = []string{
132 "9f12c13e52bca96ad4882a26558e48ab67ddd63e062b839207e893d961390f01",
133 "16d246dd6826ec7346c7328f11c4261facf82d4689f33263ff6e207956a77f21",
134 "4a22cc901c6337daa17a431fa20170684b710e5f551509511492ec24e81a8f2f",
135 "470d61abcbd154977bffc9a9ec5a8daff0caabcf2a25e8441f604c79daa0f82d",
136 }
137 for i := 0; i < len(keys); i++ {
138 c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
139 }
140
141 keys, err = tree.getKeysAtLevel(4)
142 c.Assert(err, qt.IsNil)
143 expected = []string{
144 "7a5d1c81f7b96318012de3417e53d4f13df5b1337718651cd29d0cb0a66edd20",
145 "3408213e4e844bdf3355eb8781c74e31626812898c2dbe141ed6d2c92256fc1c",
146 "dfd8a4d0b6954a3e9f3892e655b58d456eeedf9367f27dfdd9bc2dd6a5577312",
147 "9e99fbec06fb2a6725997c12c4995f62725eb4cce4808523a5a5e80cca64b007",
148 "0befa1e070231dbf4e8ff841c05878cdec823e0c09594c24910a248b3ff5a628",
149 "b7131b0a15c772a57005a4dc5d0d6dd4b3414f5d9ee7408ce5e86c5ab3520e04",
150 "6d1abe0364077846a56bab1deb1a04883eb796b74fe531a7676a9a370f83ab21",
151 "4270116394bede69cf9cd72069eca018238080380bef5de75be8dcbbe968e105",
152 }
153 for i := 0; i < len(keys); i++ {
154 c.Assert(hex.EncodeToString(keys[i]), qt.Equals, expected[i])
155 }
156}
157
158func TestSplitInBuckets(t *testing.T) {
159 c := qt.New(t)
160
161 nLeafs := 16
162 kvs := make([]kv, nLeafs)
163 for i := 0; i < nLeafs; i++ {
164 k := BigIntToBytes(big.NewInt(int64(i)))
165 v := BigIntToBytes(big.NewInt(int64(i * 2)))
166 keyPath := make([]byte, 32)
167 copy(keyPath[:], k)
168 kvs[i].pos = i
169 kvs[i].keyPath = k
170 kvs[i].k = k
171 kvs[i].v = v
172 }
173
174 // check keyToBucket results for 4 buckets & 8 keys
175 c.Assert(keyToBucket(kvs[0].k, 4), qt.Equals, 0)
176 c.Assert(keyToBucket(kvs[1].k, 4), qt.Equals, 2)
177 c.Assert(keyToBucket(kvs[2].k, 4), qt.Equals, 1)
178 c.Assert(keyToBucket(kvs[3].k, 4), qt.Equals, 3)
179 c.Assert(keyToBucket(kvs[4].k, 4), qt.Equals, 0)
180 c.Assert(keyToBucket(kvs[5].k, 4), qt.Equals, 2)
181 c.Assert(keyToBucket(kvs[6].k, 4), qt.Equals, 1)
182 c.Assert(keyToBucket(kvs[7].k, 4), qt.Equals, 3)
183
184 // check keyToBucket results for 8 buckets & 8 keys
185 c.Assert(keyToBucket(kvs[0].k, 8), qt.Equals, 0)
186 c.Assert(keyToBucket(kvs[1].k, 8), qt.Equals, 4)
187 c.Assert(keyToBucket(kvs[2].k, 8), qt.Equals, 2)
188 c.Assert(keyToBucket(kvs[3].k, 8), qt.Equals, 6)
189 c.Assert(keyToBucket(kvs[4].k, 8), qt.Equals, 1)
190 c.Assert(keyToBucket(kvs[5].k, 8), qt.Equals, 5)
191 c.Assert(keyToBucket(kvs[6].k, 8), qt.Equals, 3)
192 c.Assert(keyToBucket(kvs[7].k, 8), qt.Equals, 7)
193
194 buckets := splitInBuckets(kvs, 4)
195
196 expected := [][]string{
197 {
198 "00000000", // bucket 0
199 "08000000",
200 "04000000",
201 "0c000000",
202 },
203 {
204 "02000000", // bucket 1
205 "0a000000",
206 "06000000",
207 "0e000000",
208 },
209 {
210 "01000000", // bucket 2
211 "09000000",
212 "05000000",
213 "0d000000",
214 },
215 {
216 "03000000", // bucket 3
217 "0b000000",
218 "07000000",
219 "0f000000",
220 },
221 }
222
223 for i := 0; i < len(buckets); i++ {
224 sortKvs(buckets[i])
225 c.Assert(len(buckets[i]), qt.Equals, len(expected[i]))
226 for j := 0; j < len(buckets[i]); j++ {
227 c.Check(hex.EncodeToString(buckets[i][j].k[:4]), qt.Equals, expected[i][j])
228 }
229 }
230}
231
232func TestAddBatchCaseC(t *testing.T) {
233 c := qt.New(t)
234
235 nLeafs := 1024
236
237 tree, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
238 c.Assert(err, qt.IsNil)
239 defer tree.db.Close()
240
241 start := time.Now()
242 for i := 0; i < nLeafs; i++ {
243 k := BigIntToBytes(big.NewInt(int64(i)))
244 v := BigIntToBytes(big.NewInt(int64(i * 2)))
245 if err := tree.Add(k, v); err != nil {
246 t.Fatal(err)
247 }
248 }
249 fmt.Println(time.Since(start))
250
251 tree2, err := NewTree(memory.NewMemoryStorage(), 100, HashFunctionPoseidon)
252 c.Assert(err, qt.IsNil)
253 defer tree2.db.Close()
254
255 // add the initial leafs to fill a bit the tree before calling the
256 // AddBatch method
257 for i := 0; i < 101; i++ { // TMP TODO use const minLeafsThreshold-1 once ready
258 k := BigIntToBytes(big.NewInt(int64(i)))
259 v := BigIntToBytes(big.NewInt(int64(i * 2)))
260 if err := tree2.Add(k, v); err != nil {
261 t.Fatal(err)
262 }
263 }
264 // tree2.PrintGraphviz(nil)
265
266 var keys, values [][]byte
267 for i := 101; i < nLeafs; i++ {
268 k := BigIntToBytes(big.NewInt(int64(i)))
269 v := BigIntToBytes(big.NewInt(int64(i * 2)))
270 keys = append(keys, k)
271 values = append(values, v)
272 }
273 start = time.Now()
274 indexes, err := tree2.AddBatchOpt(keys, values)
275 c.Assert(err, qt.IsNil)
276 fmt.Println(time.Since(start))
277 c.Check(len(indexes), qt.Equals, 0)
278
279 // check that both trees roots are equal
280 c.Check(tree2.Root(), qt.DeepEquals, tree.Root())
281
282 // tree.PrintGraphviz(nil)
283 // tree2.PrintGraphviz(nil)
284 // // tree.PrintGraphvizFirstNLevels(nil, 4)
285 // // tree2.PrintGraphvizFirstNLevels(nil, 4)
286 // fmt.Println("TREE")
287 // printLeafs("t1.txt", tree)
288 // fmt.Println("TREE2")
289 // printLeafs("t2.txt", tree2)
290}
291
292// func printLeafs(name string, t *Tree) {
293// w := bytes.NewBufferString("")
294//
295// err := t.Iterate(func(k, v []byte) {
296// if v[0] != PrefixValueLeaf {
297// return
298// }
299// leafK, _ := readLeafValue(v)
300// fmt.Fprintf(w, hex.EncodeToString(leafK[:4])+"\n")
301// })
302// if err != nil {
303// panic(err)
304// }
305// err = ioutil.WriteFile(name, w.Bytes(), 0644)
306// if err != nil {
307// panic(err)
308// }
309//
310// }
311
312// func TestComputeCosts(t *testing.T) {
313// fmt.Println(computeSimpleAddCost(10))
314// fmt.Println(computeBottomUpAddCost(10))
315//
316// fmt.Println(computeSimpleAddCost(1024))
317// fmt.Println(computeBottomUpAddCost(1024))
318// }
319
320// TODO test tree with nLeafs > minLeafsThreshold, but that at level L, there is
321// less keys than nBuckets (so CASE C could be applied if first few leafs are
322// added to balance the tree)