asmt/trie_test.go
trie_test.goRaw
1/**
2 * @file
3 * @copyright defined in aergo/LICENSE.txt
4 */
5
6package trie
7
8import (
9 "bytes"
10 "runtime"
11
12 //"io/ioutil"
13 "os"
14 "path"
15 "time"
16
17 //"encoding/hex"
18 "fmt"
19 "math/rand"
20 "sort"
21 "testing"
22
23 "github.com/aergoio/aergo-lib/db"
24)
25
26func TestTrieEmpty(t *testing.T) {
27 smt := NewTrie(nil, Hasher, nil)
28 if len(smt.Root) != 0 {
29 t.Fatal("empty trie root hash not correct")
30 }
31}
32
33func TestTrieUpdateAndGet(t *testing.T) {
34 smt := NewTrie(nil, Hasher, nil)
35 smt.atomicUpdate = false
36
37 // Add data to empty trie
38 keys := getFreshData(10, 32)
39 values := getFreshData(10, 32)
40 ch := make(chan mresult, 1)
41 smt.update(smt.Root, keys, values, nil, 0, smt.TrieHeight, ch)
42 res := <-ch
43 root := res.update
44
45 // Check all keys have been stored
46 for i, key := range keys {
47 value, _ := smt.get(root, key, nil, 0, smt.TrieHeight)
48 if !bytes.Equal(values[i], value) {
49 t.Fatal("value not updated")
50 }
51 }
52
53 // Append to the trie
54 newKeys := getFreshData(5, 32)
55 newValues := getFreshData(5, 32)
56 ch = make(chan mresult, 1)
57 smt.update(root, newKeys, newValues, nil, 0, smt.TrieHeight, ch)
58 res = <-ch
59 newRoot := res.update
60 if bytes.Equal(root, newRoot) {
61 t.Fatal("trie not updated")
62 }
63 for i, newKey := range newKeys {
64 newValue, _ := smt.get(newRoot, newKey, nil, 0, smt.TrieHeight)
65 if !bytes.Equal(newValues[i], newValue) {
66 t.Fatal("failed to get value")
67 }
68 }
69}
70
71func TestTrieAtomicUpdate(t *testing.T) {
72 smt := NewTrie(nil, Hasher, nil)
73 smt.CacheHeightLimit = 0
74 keys := getFreshData(1, 32)
75 values := getFreshData(1, 32)
76 root, _ := smt.AtomicUpdate(keys, values)
77 updatedNb := len(smt.db.updatedNodes)
78 cacheNb := len(smt.db.liveCache)
79 newvalues := getFreshData(1, 32)
80 smt.AtomicUpdate(keys, newvalues)
81 if len(smt.db.updatedNodes) != 2*updatedNb {
82 t.Fatal("Atomic update doesnt store all tries")
83 }
84 if len(smt.db.liveCache) != cacheNb {
85 t.Fatal("Cache size should remain the same")
86 }
87
88 // check keys of previous atomic update are accessible in
89 // updated nodes with root.
90 smt.atomicUpdate = false
91 for i, key := range keys {
92 value, _ := smt.get(root, key, nil, 0, smt.TrieHeight)
93 if !bytes.Equal(values[i], value) {
94 t.Fatal("failed to get value")
95 }
96 }
97}
98
99func TestTriePublicUpdateAndGet(t *testing.T) {
100 smt := NewTrie(nil, Hasher, nil)
101 smt.CacheHeightLimit = 0
102 // Add data to empty trie
103 keys := getFreshData(20, 32)
104 values := getFreshData(20, 32)
105 root, _ := smt.Update(keys, values)
106 updatedNb := len(smt.db.updatedNodes)
107 cacheNb := len(smt.db.liveCache)
108
109 // Check all keys have been stored
110 for i, key := range keys {
111 value, _ := smt.Get(key)
112 if !bytes.Equal(values[i], value) {
113 t.Fatal("trie not updated")
114 }
115 }
116 if !bytes.Equal(root, smt.Root) {
117 t.Fatal("Root not stored")
118 }
119
120 newValues := getFreshData(20, 32)
121 smt.Update(keys, newValues)
122
123 if len(smt.db.updatedNodes) != updatedNb {
124 t.Fatal("multiple updates don't actualise updated nodes")
125 }
126 if len(smt.db.liveCache) != cacheNb {
127 t.Fatal("multiple updates don't actualise liveCache")
128 }
129 // Check all keys have been modified
130 for i, key := range keys {
131 value, _ := smt.Get(key)
132 if !bytes.Equal(newValues[i], value) {
133 t.Fatal("trie not updated")
134 }
135 }
136}
137
138func TestGetWithRoot(t *testing.T) {
139 dbPath := t.TempDir()
140 st := db.NewDB(db.BadgerImpl, dbPath)
141 smt := NewTrie(nil, Hasher, st)
142 smt.CacheHeightLimit = 0
143
144 // Add data to empty trie
145 keys := getFreshData(20, 32)
146 values := getFreshData(20, 32)
147 root, _ := smt.Update(keys, values)
148
149 // Check all keys have been stored
150 for i, key := range keys {
151 value, _ := smt.Get(key)
152 if !bytes.Equal(values[i], value) {
153 t.Fatal("trie not updated")
154 }
155 }
156 if !bytes.Equal(root, smt.Root) {
157 t.Fatal("Root not stored")
158 }
159 if err := smt.Commit(); err != nil {
160 t.Fatal(err)
161 }
162 // Delete two values (0 and 1)
163 if _, err := smt.Update([][]byte{keys[0], keys[1]}, [][]byte{DefaultLeaf, DefaultLeaf}); err != nil {
164 t.Fatal(err)
165 }
166
167 // Change one value
168 oldValue3 := make([]byte, 32)
169 copy(oldValue3, values[3])
170 values[3] = getFreshData(1, 32)[0]
171 if _, err := smt.Update([][]byte{keys[3]}, [][]byte{values[3]}); err != nil {
172 t.Fatal(err)
173 }
174
175 // Check root has been actually updated
176 if bytes.Equal(smt.Root, root) {
177 t.Fatal("root not updated")
178 }
179
180 // Get the third value with the new root
181 v3, err := smt.GetWithRoot(keys[3], smt.Root)
182 if err != nil {
183 t.Fatal(err)
184 }
185 if !bytes.Equal(v3, values[3]) {
186 t.Fatalf("GetWithRoot did not keep the value: %x != %x", v3, values[3])
187 }
188
189 // Get the third value with the old root
190 v3, err = smt.GetWithRoot(keys[3], root)
191 if err != nil {
192 t.Fatal(err)
193 }
194 if !bytes.Equal(v3, oldValue3) {
195 t.Fatalf("GetWithRoot did not keep the value: %x != %x", v3, oldValue3)
196 }
197 st.Close()
198}
199
200func TestTrieWalk(t *testing.T) {
201 dbPath := t.TempDir()
202 st := db.NewDB(db.BadgerImpl, dbPath)
203
204 smt := NewTrie(nil, Hasher, st)
205 smt.CacheHeightLimit = 0
206 // Add data to empty trie
207 keys := getFreshData(20, 32)
208 values := getFreshData(20, 32)
209 root, _ := smt.Update(keys, values)
210
211 // Check all keys have been stored
212 for i, key := range keys {
213 value, _ := smt.Get(key)
214 if !bytes.Equal(values[i], value) {
215 t.Fatal("trie not updated")
216 }
217 }
218 if !bytes.Equal(root, smt.Root) {
219 t.Fatal("Root not stored")
220 }
221
222 // Walk over the whole tree and compare the values
223 i := 0
224 if err := smt.Walk(nil, func(v *WalkResult) int32 {
225 if !bytes.Equal(v.Value, values[i]) {
226 t.Fatalf("walk value does not match %x != %x", v.Value, values[i])
227 }
228 if !bytes.Equal(v.Key, keys[i]) {
229 t.Fatalf("walk key does not match %x != %x", v.Key, keys[i])
230 }
231 i++
232 return 0
233 }); err != nil {
234 t.Fatal(err)
235 }
236
237 // Delete two values (0 and 3)
238 if _, err := smt.Update([][]byte{keys[0], keys[3]}, [][]byte{DefaultLeaf, DefaultLeaf}); err != nil {
239 t.Fatal(err)
240 }
241
242 // Delete two elements and walk again
243 i = 1
244 if err := smt.Walk(nil, func(v *WalkResult) int32 {
245 if i == 3 {
246 i++
247 }
248 if !bytes.Equal(v.Value, values[i]) {
249 t.Fatalf("walk value does not match %x == %x\n", v.Value, values[i])
250 }
251 if !bytes.Equal(v.Key, keys[i]) {
252 t.Fatalf("walk key does not match %x == %x\n", v.Key, keys[i])
253 }
254 i++
255 return 0
256 }); err != nil {
257 t.Fatal(err)
258 }
259
260 // Add one new value to preivous deleted key
261 values[3] = getFreshData(1, 32)[0]
262 if _, err := smt.Update([][]byte{keys[3]}, [][]byte{values[3]}); err != nil {
263 t.Fatal(err)
264 }
265
266 // Walk and check again
267 i = 1
268 if err := smt.Walk(nil, func(v *WalkResult) int32 {
269 if !bytes.Equal(v.Value, values[i]) {
270 t.Fatalf("walk value does not match %x != %x\n", v.Value, values[i])
271 }
272 if !bytes.Equal(v.Key, keys[i]) {
273 t.Fatalf("walk key does not match %x != %x\n", v.Key, keys[i])
274 }
275 i++
276 return 0
277 }); err != nil {
278 t.Fatal(err)
279 }
280
281 // Find a specific value and test stop
282 i = 0
283 if err := smt.Walk(nil, func(v *WalkResult) int32 {
284 if bytes.Equal(v.Value, values[5]) {
285 return 1
286 }
287 i++
288 return 0
289 }); err != nil {
290 t.Fatal(err)
291 }
292 if i != 4 {
293 t.Fatalf("Needed more iterations on walk than expected: %d != 4", i)
294 }
295
296 st.Close()
297}
298
299func TestTrieDelete(t *testing.T) {
300 smt := NewTrie(nil, Hasher, nil)
301 // Add data to empty trie
302 keys := getFreshData(20, 32)
303 values := getFreshData(20, 32)
304 ch := make(chan mresult, 1)
305 smt.update(smt.Root, keys, values, nil, 0, smt.TrieHeight, ch)
306 result := <-ch
307 root := result.update
308 value, _ := smt.get(root, keys[0], nil, 0, smt.TrieHeight)
309 if !bytes.Equal(values[0], value) {
310 t.Fatal("trie not updated")
311 }
312
313 // Delete from trie
314 // To delete a key, just set it's value to Default leaf hash.
315 ch = make(chan mresult, 1)
316 smt.update(root, keys[0:1], [][]byte{DefaultLeaf}, nil, 0, smt.TrieHeight, ch)
317 result = <-ch
318 updatedNb := len(smt.db.updatedNodes)
319 newRoot := result.update
320 newValue, _ := smt.get(newRoot, keys[0], nil, 0, smt.TrieHeight)
321 if len(newValue) != 0 {
322 t.Fatal("Failed to delete from trie")
323 }
324 // Remove deleted key from keys and check root with a clean trie.
325 smt2 := NewTrie(nil, Hasher, nil)
326 ch = make(chan mresult, 1)
327 smt2.update(smt.Root, keys[1:], values[1:], nil, 0, smt.TrieHeight, ch)
328 result = <-ch
329 cleanRoot := result.update
330 if !bytes.Equal(newRoot, cleanRoot) {
331 t.Fatal("roots mismatch")
332 }
333
334 if len(smt2.db.updatedNodes) != updatedNb {
335 t.Fatal("deleting doesn't actualise updated nodes")
336 }
337
338 //Empty the trie
339 var newValues [][]byte
340 for i := 0; i < 20; i++ {
341 newValues = append(newValues, DefaultLeaf)
342 }
343 ch = make(chan mresult, 1)
344 smt.update(root, keys, newValues, nil, 0, smt.TrieHeight, ch)
345 result = <-ch
346 root = result.update
347 //if !bytes.Equal(smt.DefaultHash(256), root) {
348 if len(root) != 0 {
349 t.Fatal("empty trie root hash not correct")
350 }
351
352 // Test deleting an already empty key
353 smt = NewTrie(nil, Hasher, nil)
354 keys = getFreshData(2, 32)
355 values = getFreshData(2, 32)
356 root, _ = smt.Update(keys, values)
357 key0 := make([]byte, 32, 32)
358 key1 := make([]byte, 32, 32)
359 smt.Update([][]byte{key0, key1}, [][]byte{DefaultLeaf, DefaultLeaf})
360 if !bytes.Equal(root, smt.Root) {
361 t.Fatal("deleting a default key shouldnt' modify the tree")
362 }
363}
364
365// test updating and deleting at the same time
366func TestTrieUpdateAndDelete(t *testing.T) {
367 smt := NewTrie(nil, Hasher, nil)
368 smt.CacheHeightLimit = 0
369 key0 := make([]byte, 32, 32)
370 values := getFreshData(1, 32)
371 root, _ := smt.Update([][]byte{key0}, values)
372 cacheNb := len(smt.db.liveCache)
373 updatedNb := len(smt.db.updatedNodes)
374 smt.atomicUpdate = false
375 _, _, k, v, isShortcut, _ := smt.loadChildren(root, smt.TrieHeight, 0, nil)
376 if !isShortcut || !bytes.Equal(k[:HashLength], key0) || !bytes.Equal(v[:HashLength], values[0]) {
377 t.Fatal("leaf shortcut didn't move up to root")
378 }
379
380 key1 := make([]byte, 32, 32)
381 // set the last bit
382 bitSet(key1, 255)
383 keys := [][]byte{key0, key1}
384 values = [][]byte{DefaultLeaf, getFreshData(1, 32)[0]}
385 root, _ = smt.Update(keys, values)
386
387 if len(smt.db.liveCache) != cacheNb {
388 t.Fatal("number of cache nodes not correct after delete")
389 }
390 if len(smt.db.updatedNodes) != updatedNb {
391 t.Fatal("number of cache nodes not correct after delete")
392 }
393
394 smt.atomicUpdate = false
395 _, _, k, v, isShortcut, _ = smt.loadChildren(root, smt.TrieHeight, 0, nil)
396 if !isShortcut || !bytes.Equal(k[:HashLength], key1) || !bytes.Equal(v[:HashLength], values[1]) {
397 t.Fatal("leaf shortcut didn't move up to root")
398 }
399}
400
401func TestTrieMerkleProof(t *testing.T) {
402 smt := NewTrie(nil, Hasher, nil)
403 // Add data to empty trie
404 keys := getFreshData(10, 32)
405 values := getFreshData(10, 32)
406 smt.Update(keys, values)
407
408 for i, key := range keys {
409 ap, _, k, v, _ := smt.MerkleProof(key)
410 if !smt.VerifyInclusion(ap, key, values[i]) {
411 t.Fatalf("failed to verify inclusion proof")
412 }
413 if !bytes.Equal(key, k) && !bytes.Equal(values[i], v) {
414 t.Fatalf("merkle proof didnt return the correct key-value pair")
415 }
416 }
417 emptyKey := Hasher([]byte("non-member"))
418 ap, included, proofKey, proofValue, _ := smt.MerkleProof(emptyKey)
419 if included {
420 t.Fatalf("failed to verify non inclusion proof")
421 }
422 if !smt.VerifyNonInclusion(ap, emptyKey, proofValue, proofKey) {
423 t.Fatalf("failed to verify non inclusion proof")
424 }
425}
426
427func TestTrieMerkleProofCompressed(t *testing.T) {
428 smt := NewTrie(nil, Hasher, nil)
429 // Add data to empty trie
430 keys := getFreshData(10, 32)
431 values := getFreshData(10, 32)
432 smt.Update(keys, values)
433
434 for i, key := range keys {
435 bitmap, ap, length, _, k, v, _ := smt.MerkleProofCompressed(key)
436 if !smt.VerifyInclusionC(bitmap, key, values[i], ap, length) {
437 t.Fatalf("failed to verify inclusion proof")
438 }
439 if !bytes.Equal(key, k) && !bytes.Equal(values[i], v) {
440 t.Fatalf("merkle proof didnt return the correct key-value pair")
441 }
442 }
443 emptyKey := Hasher([]byte("non-member"))
444 bitmap, ap, length, included, proofKey, proofValue, _ := smt.MerkleProofCompressed(emptyKey)
445 if included {
446 t.Fatalf("failed to verify non inclusion proof")
447 }
448 if !smt.VerifyNonInclusionC(ap, length, bitmap, emptyKey, proofValue, proofKey) {
449 t.Fatalf("failed to verify non inclusion proof")
450 }
451}
452
453func TestTrieCommit(t *testing.T) {
454 dbPath := path.Join(".aergo", "db")
455 if _, err := os.Stat(dbPath); os.IsNotExist(err) {
456 _ = os.MkdirAll(dbPath, 0711)
457 }
458 st := db.NewDB(db.BadgerImpl, dbPath)
459
460 smt := NewTrie(nil, Hasher, st)
461 keys := getFreshData(10, 32)
462 values := getFreshData(10, 32)
463 smt.Update(keys, values)
464 smt.Commit()
465 // liveCache is deleted so the key is fetched in badger db
466 smt.db.liveCache = make(map[Hash][][]byte)
467 for i, key := range keys {
468 value, _ := smt.Get(key)
469 if !bytes.Equal(value, values[i]) {
470 t.Fatal("failed to get value in committed db")
471 }
472 }
473 st.Close()
474 os.RemoveAll(".aergo")
475}
476
477func TestTrieStageUpdates(t *testing.T) {
478 dbPath := path.Join(".aergo", "db")
479 if _, err := os.Stat(dbPath); os.IsNotExist(err) {
480 _ = os.MkdirAll(dbPath, 0711)
481 }
482 st := db.NewDB(db.BadgerImpl, dbPath)
483
484 smt := NewTrie(nil, Hasher, st)
485 keys := getFreshData(10, 32)
486 values := getFreshData(10, 32)
487 smt.Update(keys, values)
488 txn := st.NewTx()
489 smt.StageUpdates(txn.(DbTx))
490 txn.Commit()
491 // liveCache is deleted so the key is fetched in badger db
492 smt.db.liveCache = make(map[Hash][][]byte)
493 for i, key := range keys {
494 value, _ := smt.Get(key)
495 if !bytes.Equal(value, values[i]) {
496 t.Fatal("failed to get value in committed db")
497 }
498 }
499 st.Close()
500 os.RemoveAll(".aergo")
501}
502
503func TestTrieRevert(t *testing.T) {
504 dbPath := path.Join(".aergo", "db")
505 if _, err := os.Stat(dbPath); os.IsNotExist(err) {
506 _ = os.MkdirAll(dbPath, 0711)
507 }
508 st := db.NewDB(db.BadgerImpl, dbPath)
509
510 smt := NewTrie(nil, Hasher, st)
511
512 // Edge case : Test that revert doesnt delete shortcut nodes
513 // when moved to a different position in tree
514 key0 := make([]byte, 32, 32)
515 key1 := make([]byte, 32, 32)
516 // setting the bit at 251 creates 2 shortcut batches at height 252
517 bitSet(key1, 251)
518 values := getFreshData(2, 32)
519 keys := [][]byte{key0, key1}
520 root, _ := smt.Update([][]byte{key0}, [][]byte{values[0]})
521 smt.Commit()
522 root2, _ := smt.Update([][]byte{key1}, [][]byte{values[1]})
523 smt.Commit()
524 smt.Revert(root)
525 if len(smt.db.Store.Get(root)) == 0 {
526 t.Fatal("shortcut node shouldnt be deleted by revert")
527 }
528 if len(smt.db.Store.Get(root2)) != 0 {
529 t.Fatal("reverted root should have been deleted")
530 }
531 key1 = make([]byte, 32, 32)
532 // setting the bit at 255 stores the keys as the tip
533 bitSet(key1, 255)
534 smt.Update([][]byte{key1}, [][]byte{values[1]})
535 smt.Commit()
536 smt.Revert(root)
537 if len(smt.db.Store.Get(root)) == 0 {
538 t.Fatal("shortcut node shouldnt be deleted by revert")
539 }
540
541 // Test all nodes are reverted in the usual case
542 // Add data to empty trie
543 keys = getFreshData(10, 32)
544 values = getFreshData(10, 32)
545 root, _ = smt.Update(keys, values)
546 smt.Commit()
547
548 // Update the values
549 newValues := getFreshData(10, 32)
550 smt.Update(keys, newValues)
551 updatedNodes1 := smt.db.updatedNodes
552 smt.Commit()
553 newKeys := getFreshData(10, 32)
554 newValues = getFreshData(10, 32)
555 smt.Update(newKeys, newValues)
556 updatedNodes2 := smt.db.updatedNodes
557 smt.Commit()
558
559 smt.Revert(root)
560
561 if !bytes.Equal(smt.Root, root) {
562 t.Fatal("revert failed")
563 }
564 if len(smt.pastTries) != 2 { // contains empty trie + reverted trie
565 t.Fatal("past tries not updated after revert")
566 }
567 // Check all keys have been reverted
568 for i, key := range keys {
569 value, _ := smt.Get(key)
570 if !bytes.Equal(values[i], value) {
571 t.Fatal("revert failed, values not updated")
572 }
573 }
574 if len(smt.db.liveCache) != 0 {
575 t.Fatal("live cache not reset after revert")
576 }
577 // Check all reverted nodes have been deleted
578 for node, _ := range updatedNodes2 {
579 if len(smt.db.Store.Get(node[:])) != 0 {
580 t.Fatal("nodes not deleted from database", node)
581 }
582 }
583 for node, _ := range updatedNodes1 {
584 if len(smt.db.Store.Get(node[:])) != 0 {
585 t.Fatal("nodes not deleted from database", node)
586 }
587 }
588 st.Close()
589 os.RemoveAll(".aergo")
590}
591
592func TestTrieRaisesError(t *testing.T) {
593 dbPath := path.Join(".aergo", "db")
594 if _, err := os.Stat(dbPath); os.IsNotExist(err) {
595 _ = os.MkdirAll(dbPath, 0711)
596 }
597 st := db.NewDB(db.BadgerImpl, dbPath)
598
599 smt := NewTrie(nil, Hasher, st)
600 // Add data to empty trie
601 keys := getFreshData(10, 32)
602 values := getFreshData(10, 32)
603 smt.Update(keys, values)
604 smt.db.liveCache = make(map[Hash][][]byte)
605 smt.db.updatedNodes = make(map[Hash][][]byte)
606
607 // Check errors are raised is a keys is not in cache nore db
608 for _, key := range keys {
609 _, err := smt.Get(key)
610 if err == nil {
611 t.Fatal("Error not created if database doesnt have a node")
612 }
613 }
614 _, _, _, _, _, _, err := smt.MerkleProofCompressed(keys[0])
615 if err == nil {
616 t.Fatal("Error not created if database doesnt have a node")
617 }
618 _, err = smt.Update(keys, values)
619 if err == nil {
620 t.Fatal("Error not created if database doesnt have a node")
621 }
622 st.Close()
623 os.RemoveAll(".aergo")
624
625 smt = NewTrie(nil, Hasher, nil)
626 err = smt.Commit()
627 if err == nil {
628 t.Fatal("Error not created if database not connected")
629 }
630 smt.db.liveCache = make(map[Hash][][]byte)
631 smt.atomicUpdate = false
632 _, _, _, _, _, err = smt.loadChildren(make([]byte, 32, 32), smt.TrieHeight, 0, nil)
633 if err == nil {
634 t.Fatal("Error not created if database not connected")
635 }
636 err = smt.LoadCache(make([]byte, 32))
637 if err == nil {
638 t.Fatal("Error not created if database not connected")
639 }
640}
641
642func TestTrieLoadCache(t *testing.T) {
643 dbPath := path.Join(".aergo", "db")
644 if _, err := os.Stat(dbPath); os.IsNotExist(err) {
645 _ = os.MkdirAll(dbPath, 0711)
646 }
647 st := db.NewDB(db.BadgerImpl, dbPath)
648
649 smt := NewTrie(nil, Hasher, st)
650 // Test size of cache
651 smt.CacheHeightLimit = 0
652 key0 := make([]byte, 32, 32)
653 key1 := make([]byte, 32, 32)
654 bitSet(key1, 255)
655 values := getFreshData(2, 32)
656 smt.Update([][]byte{key0, key1}, values)
657 if len(smt.db.liveCache) != 66 {
658 // the nodes are at the tip, so 64 + 2 = 66
659 t.Fatal("cache size incorrect")
660 }
661
662 // Add data to empty trie
663 keys := getFreshData(10, 32)
664 values = getFreshData(10, 32)
665 smt.Update(keys, values)
666 smt.Commit()
667
668 // Simulate node restart by deleting and loading cache
669 cacheSize := len(smt.db.liveCache)
670 smt.db.liveCache = make(map[Hash][][]byte)
671
672 err := smt.LoadCache(smt.Root)
673
674 if err != nil {
675 t.Fatal(err)
676 }
677 if cacheSize != len(smt.db.liveCache) {
678 t.Fatal("Cache loading from db incorrect")
679 }
680 st.Close()
681 os.RemoveAll(".aergo")
682}
683
684func TestHeight0LeafShortcut(t *testing.T) {
685 keySize := 32
686 smt := NewTrie(nil, Hasher, nil)
687 // Add 2 sibling keys that will be stored at height 0
688 key0 := make([]byte, keySize, keySize)
689 key1 := make([]byte, keySize, keySize)
690 bitSet(key1, keySize*8-1)
691 keys := [][]byte{key0, key1}
692 values := getFreshData(2, 32)
693 smt.Update(keys, values)
694 updatedNb := len(smt.db.updatedNodes)
695
696 // Check all keys have been stored
697 for i, key := range keys {
698 value, _ := smt.Get(key)
699 if !bytes.Equal(values[i], value) {
700 t.Fatal("trie not updated")
701 }
702 }
703 bitmap, ap, length, _, k, v, err := smt.MerkleProofCompressed(key1)
704 if err != nil {
705 t.Fatal(err)
706 }
707 if !bytes.Equal(key1, k) && !bytes.Equal(values[1], v) {
708 t.Fatalf("merkle proof didnt return the correct key-value pair")
709 }
710 if length != smt.TrieHeight {
711 t.Fatal("proof should have length equal to trie height for a leaf shortcut")
712 }
713 if !smt.VerifyInclusionC(bitmap, key1, values[1], ap, length) {
714 t.Fatal("failed to verify inclusion proof")
715 }
716
717 // Delete one key and check that the remaining one moved up to the root of the tree
718 newRoot, _ := smt.AtomicUpdate(keys[0:1], [][]byte{DefaultLeaf})
719
720 // Nb of updated nodes remains same because the new shortcut root was already stored at height 0.
721 if len(smt.db.updatedNodes) != updatedNb {
722 fmt.Println(len(smt.db.updatedNodes), updatedNb)
723 t.Fatal("number of cache nodes not correct after delete")
724 }
725 smt.atomicUpdate = false
726 _, _, k, v, isShortcut, err := smt.loadChildren(newRoot, smt.TrieHeight, 0, nil)
727 if err != nil {
728 t.Fatal(err)
729 }
730 if !isShortcut || !bytes.Equal(k[:HashLength], key1) || !bytes.Equal(v[:HashLength], values[1]) {
731 t.Fatal("leaf shortcut didn't move up to root")
732 }
733
734 _, _, length, _, k, v, _ = smt.MerkleProofCompressed(key1)
735 if length != 0 {
736 t.Fatal("proof should have length equal to trie height for a leaf shortcut")
737 }
738 if !bytes.Equal(key1, k) && !bytes.Equal(values[1], v) {
739 t.Fatalf("merkle proof didnt return the correct key-value pair")
740 }
741}
742
743func TestStash(t *testing.T) {
744 dbPath := path.Join(".aergo", "db")
745 if _, err := os.Stat(dbPath); os.IsNotExist(err) {
746 _ = os.MkdirAll(dbPath, 0711)
747 }
748 st := db.NewDB(db.BadgerImpl, dbPath)
749 smt := NewTrie(nil, Hasher, st)
750 // Add data to empty trie
751 keys := getFreshData(20, 32)
752 values := getFreshData(20, 32)
753 root, _ := smt.Update(keys, values)
754 cacheSize := len(smt.db.liveCache)
755 smt.Commit()
756 if len(smt.pastTries) != 1 {
757 t.Fatal("Past tries not updated after commit")
758 }
759 values = getFreshData(20, 32)
760 smt.Update(keys, values)
761 smt.Stash(true)
762 if len(smt.pastTries) != 1 {
763 t.Fatal("Past tries not updated after commit")
764 }
765 if !bytes.Equal(smt.Root, root) {
766 t.Fatal("Trie not rolled back")
767 }
768 if len(smt.db.updatedNodes) != 0 {
769 t.Fatal("Trie not rolled back")
770 }
771 if len(smt.db.liveCache) != cacheSize {
772 t.Fatal("Trie not rolled back")
773 }
774 keys = getFreshData(20, 32)
775 values = getFreshData(20, 32)
776 smt.AtomicUpdate(keys, values)
777 values = getFreshData(20, 32)
778 smt.AtomicUpdate(keys, values)
779 if len(smt.pastTries) != 3 {
780 t.Fatal("Past tries not updated after commit")
781 }
782 smt.Stash(true)
783 if !bytes.Equal(smt.Root, root) {
784 t.Fatal("Trie not rolled back")
785 }
786 if len(smt.db.updatedNodes) != 0 {
787 t.Fatal("Trie not rolled back")
788 }
789 if len(smt.db.liveCache) != cacheSize {
790 t.Fatal("Trie not rolled back")
791 }
792 if len(smt.pastTries) != 1 {
793 t.Fatal("Past tries not updated after commit")
794 }
795 st.Close()
796 os.RemoveAll(".aergo")
797}
798
799func benchmark10MAccounts10Ktps(smt *Trie, b *testing.B) {
800 //b.ReportAllocs()
801 keys := getFreshData(100, 32)
802 values := getFreshData(100, 32)
803 smt.Update(keys, values)
804 fmt.Println("\nLoading b.N x 1000 accounts")
805 for i := 0; i < b.N; i++ {
806 newkeys := getFreshData(1000, 32)
807 newvalues := getFreshData(1000, 32)
808 start := time.Now()
809 smt.Update(newkeys, newvalues)
810 end := time.Now()
811 smt.Commit()
812 end2 := time.Now()
813 for j, key := range newkeys {
814 val, _ := smt.Get(key)
815 if !bytes.Equal(val, newvalues[j]) {
816 b.Fatal("new key not included")
817 }
818 }
819 end3 := time.Now()
820 elapsed := end.Sub(start)
821 elapsed2 := end2.Sub(end)
822 elapsed3 := end3.Sub(end2)
823 var m runtime.MemStats
824 runtime.ReadMemStats(&m)
825 fmt.Println(i, " : update time : ", elapsed, "commit time : ", elapsed2,
826 "\n1000 Get time : ", elapsed3,
827 "\ndb read : ", smt.LoadDbCounter, " cache read : ", smt.LoadCacheCounter,
828 "\ncache size : ", len(smt.db.liveCache),
829 "\nRAM : ", m.Sys/1024/1024, " MiB")
830 }
831}
832
833//go test -run=xxx -bench=. -benchmem -test.benchtime=20s
834func BenchmarkCacheHeightLimit233(b *testing.B) {
835 dbPath := path.Join(".aergo", "db")
836 if _, err := os.Stat(dbPath); os.IsNotExist(err) {
837 _ = os.MkdirAll(dbPath, 0711)
838 }
839 st := db.NewDB(db.BadgerImpl, dbPath)
840 smt := NewTrie(nil, Hasher, st)
841 smt.CacheHeightLimit = 233
842 benchmark10MAccounts10Ktps(smt, b)
843 st.Close()
844 os.RemoveAll(".aergo")
845}
846func BenchmarkCacheHeightLimit238(b *testing.B) {
847 dbPath := path.Join(".aergo", "db")
848 if _, err := os.Stat(dbPath); os.IsNotExist(err) {
849 _ = os.MkdirAll(dbPath, 0711)
850 }
851 st := db.NewDB(db.BadgerImpl, dbPath)
852 smt := NewTrie(nil, Hasher, st)
853 smt.CacheHeightLimit = 238
854 benchmark10MAccounts10Ktps(smt, b)
855 st.Close()
856 os.RemoveAll(".aergo")
857}
858func BenchmarkCacheHeightLimit245(b *testing.B) {
859 dbPath := path.Join(".aergo", "db")
860 if _, err := os.Stat(dbPath); os.IsNotExist(err) {
861 _ = os.MkdirAll(dbPath, 0711)
862 }
863 st := db.NewDB(db.BadgerImpl, dbPath)
864 smt := NewTrie(nil, Hasher, st)
865 smt.CacheHeightLimit = 245
866 benchmark10MAccounts10Ktps(smt, b)
867 st.Close()
868 os.RemoveAll(".aergo")
869}
870
871func getFreshData(size, length int) [][]byte {
872 var data [][]byte
873 for i := 0; i < size; i++ {
874 key := make([]byte, 32)
875 _, err := rand.Read(key)
876 if err != nil {
877 panic(err)
878 }
879 data = append(data, Hasher(key)[:length])
880 }
881 sort.Sort(DataArray(data))
882 return data
883}