Refactor, clean up types, Readme

Author: roysc

README.md

@@ -19,21 +19,21 @@ import (
 )
 
 func main() {
-	// Initialise two new key-value store to store the nodes and values of the tree
+	// Initialise two new key-value store to store the nodes of the tree
+	// (Note: the tree only stores hashed values, not raw value data)
 	nodeStore := smt.NewSimpleMap()
-	valueStore := smt.NewSimpleMap()
 	// Initialise the tree
-	tree := smt.NewSparseMerkleTree(nodeStore, valueStore, sha256.New())
+	tree := smt.NewSMT(nodeStore, sha256.New())
 
 	// Update the key "foo" with the value "bar"
-	_, _ = tree.Update([]byte("foo"), []byte("bar"))
+	_ = tree.Update([]byte("foo"), []byte("bar"))
 
 	// Generate a Merkle proof for foo=bar
 	proof, _ := tree.Prove([]byte("foo"))
 	root := tree.Root() // We also need the current tree root for the proof
 
 	// Verify the Merkle proof for foo=bar
-	if smt.VerifyProof(proof, root, []byte("foo"), []byte("bar"), sha256.New()) {
+	if smt.VerifyProof(proof, root, []byte("foo"), []byte("bar"), tree.Spec()) {
 		fmt.Println("Proof verification succeeded.")
 	} else {
 		fmt.Println("Proof verification failed.")

bench_test.go

@@ -8,29 +8,29 @@ import (
 
 func BenchmarkSparseMerkleTree_Update(b *testing.B) {
 	smn, smv := NewSimpleMap(), NewSimpleMap()
-	smt := NewSparseMerkleTree(smn, smv, sha256.New())
+	smt := NewSMTWithStorage(smn, smv, sha256.New())
 
 	b.ResetTimer()
 	b.ReportAllocs()
 	for i := 0; i < b.N; i++ {
 		s := strconv.Itoa(i)
-		_, _ = smt.Update([]byte(s), []byte(s))
+		_ = smt.Update([]byte(s), []byte(s))
 	}
 }
 
 func BenchmarkSparseMerkleTree_Delete(b *testing.B) {
 	smn, smv := NewSimpleMap(), NewSimpleMap()
-	smt := NewSparseMerkleTree(smn, smv, sha256.New())
+	smt := NewSMTWithStorage(smn, smv, sha256.New())
 
 	for i := 0; i < 100000; i++ {
 		s := strconv.Itoa(i)
-		_, _ = smt.Update([]byte(s), []byte(s))
+		_ = smt.Update([]byte(s), []byte(s))
 	}
 
 	b.ResetTimer()
 	b.ReportAllocs()
 	for i := 0; i < b.N; i++ {
 		s := strconv.Itoa(i)
-		_, _ = smt.Delete([]byte(s))
+		_ = smt.Delete([]byte(s))
 	}
 }

bulk_test.go

@@ -4,29 +4,36 @@ import (
 	"bytes"
 	"crypto/sha256"
 	"math/rand"
-	"reflect"
 	"testing"
 )
 
+type opCounts struct{ ops, inserts, updates, deletes int }
+type bulkop struct{ key, val []byte }
+
 // Test all tree operations in bulk.
-func TestSparseMerkleTree(t *testing.T) {
-	for i := 0; i < 5; i++ {
+func TestBulkOperations(t *testing.T) {
+	rand.Seed(1)
+
+	cases := []opCounts{
 		// Test more inserts/updates than deletions.
-		bulkOperations(t, 200, 100, 100, 50)
-	}
-	for i := 0; i < 5; i++ {
+		{200, 100, 100, 50},
+		{1000, 100, 100, 50},
 		// Test extreme deletions.
-		bulkOperations(t, 200, 100, 100, 500)
+		{200, 100, 100, 500},
+		{1000, 100, 100, 500},
+	}
+	for _, tc := range cases {
+		bulkOperations(t, tc.ops, tc.inserts, tc.updates, tc.deletes)
 	}
 }
 
 // Test all tree operations in bulk, with specified ratio probabilities of insert, update and delete.
 func bulkOperations(t *testing.T, operations int, insert int, update int, delete int) {
 	smn, smv := NewSimpleMap(), NewSimpleMap()
-	smt := NewSparseMerkleTree(smn, smv, sha256.New())
+	smt := NewSMTWithStorage(smn, smv, sha256.New())
 
 	max := insert + update + delete
-	kv := make(map[string]string)
+	var kv []bulkop
 
 	for i := 0; i < operations; i++ {
 		n := rand.Intn(max)
@@ -39,98 +46,90 @@ func bulkOperations(t *testing.T, operations int, insert int, update int, delete
 			val := make([]byte, valLen)
 			rand.Read(val)
 
-			kv[string(key)] = string(val)
-			_, err := smt.Update(key, val)
+			err := smt.Update(key, val)
 			if err != nil {
-				t.Errorf("error: %v", err)
+				t.Fatalf("error: %v", err)
 			}
+			kv = append(kv, bulkop{key, val})
 		} else if n > insert && n < insert+update { // Update
-			keys := reflect.ValueOf(kv).MapKeys()
-			if len(keys) == 0 {
+			if len(kv) == 0 {
 				continue
 			}
-			key := []byte(keys[rand.Intn(len(keys))].Interface().(string))
-
+			ki := rand.Intn(len(kv))
 			valLen := 1 + rand.Intn(64)
 			val := make([]byte, valLen)
 			rand.Read(val)
 
-			kv[string(key)] = string(val)
-			_, err := smt.Update(key, val)
+			err := smt.Update(kv[ki].key, val)
 			if err != nil {
-				t.Errorf("error: %v", err)
+				t.Fatalf("error: %v", err)
 			}
+			kv[ki].val = val
 		} else { // Delete
-			keys := reflect.ValueOf(kv).MapKeys()
-			if len(keys) == 0 {
+			if len(kv) == 0 {
 				continue
 			}
-			key := []byte(keys[rand.Intn(len(keys))].Interface().(string))
+			ki := rand.Intn(len(kv))
 
-			kv[string(key)] = ""
-			_, err := smt.Update(key, defaultValue)
-			if err != nil {
-				t.Errorf("error: %v", err)
+			err := smt.Delete(kv[ki].key)
+			if err != nil && err != ErrKeyNotPresent {
+				t.Fatalf("error: %v", err)
 			}
+			kv[ki].val = defaultValue
 		}
-
-		bulkCheckAll(t, smt, &kv)
 	}
+	bulkCheckAll(t, smt, kv)
 }
 
-func bulkCheckAll(t *testing.T, smt *SparseMerkleTree, kv *map[string]string) {
-	for k, v := range *kv {
-		value, err := smt.Get([]byte(k))
+func bulkCheckAll(t *testing.T, smt *SMTWithStorage, kv []bulkop) {
+	for ki := range kv {
+		k, v := kv[ki].key, kv[ki].val
+
+		value, err := smt.GetValue([]byte(k))
 		if err != nil {
 			t.Errorf("error: %v", err)
 		}
 		if !bytes.Equal([]byte(v), value) {
-			t.Error("got incorrect value when bulk testing operations")
+			t.Errorf("Incorrect value (i=%d)", ki)
 		}
 
 		// Generate and verify a Merkle proof for this key.
 		proof, err := smt.Prove([]byte(k))
 		if err != nil {
 			t.Errorf("error: %v", err)
 		}
-		if !VerifyProof(proof, smt.Root(), []byte(k), []byte(v), smt.th.hasher) {
-			t.Error("Merkle proof failed to verify")
+		if !VerifyProof(proof, smt.Root(), []byte(k), []byte(v), smt.Spec()) {
+			t.Fatalf("Merkle proof failed to verify (i=%d): %v", ki, []byte(k))
 		}
-		compactProof, err := smt.ProveCompact([]byte(k))
+		compactProof, err := ProveCompact([]byte(k), smt)
 		if err != nil {
 			t.Errorf("error: %v", err)
 		}
-		if !VerifyCompactProof(compactProof, smt.Root(), []byte(k), []byte(v), smt.th.hasher) {
-			t.Error("Merkle proof failed to verify")
+		if !VerifyCompactProof(compactProof, smt.Root(), []byte(k), []byte(v), smt.Spec()) {
+			t.Fatalf("Compact Merkle proof failed to verify (i=%d): %v", ki, []byte(k))
 		}
 
-		if v == "" {
+		if v == nil {
 			continue
 		}
 
 		// Check that the key is at the correct height in the tree.
 		largestCommonPrefix := 0
-		for k2, v2 := range *kv {
-			if v2 == "" {
+		for ki2 := range kv {
+			k2, v2 := kv[ki2].key, kv[ki2].val
+			if v2 == nil {
 				continue
 			}
-			commonPrefix := countCommonPrefix(smt.th.path([]byte(k)), smt.th.path([]byte(k2)))
-			if commonPrefix != smt.depth() && commonPrefix > largestCommonPrefix {
+
+			ph := smt.Spec().ph
+			commonPrefix := countCommonPrefix(ph.Path([]byte(k)), ph.Path([]byte(k2)), 0)
+			if commonPrefix != smt.Spec().depth() && commonPrefix > largestCommonPrefix {
 				largestCommonPrefix = commonPrefix
 			}
 		}
-		sideNodes, _, _, _, err := smt.sideNodesForRoot(smt.th.path([]byte(k)), smt.Root(), false)
-		if err != nil {
-			t.Errorf("error: %v", err)
-		}
-		numSideNodes := 0
-		for _, v := range sideNodes {
-			if v != nil {
-				numSideNodes++
-			}
-		}
-		if numSideNodes != largestCommonPrefix+1 && (numSideNodes != 0 && largestCommonPrefix != 0) {
-			t.Error("leaf is at unexpected height")
+		if len(proof.SideNodes) != largestCommonPrefix+1 &&
+			(len(proof.SideNodes) != 0 && largestCommonPrefix != 0) {
+			t.Errorf("leaf is at unexpected height (ki=%d)", ki)
 		}
 	}
 }