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