Implement compute-optimised Merkle tree

README.md

@@ -1,13 +1,11 @@
 # smt
 
-A Go library that implements a Sparse Merkle tree for a key-value map.
+A Go library that implements a Sparse Merkle tree for a key-value map. The tree implements the same optimisations specified in the [Libra whitepaper](https://developers.libra.org/docs/assets/papers/the-libra-blockchain/2020-05-26.pdf), to reduce the number of hash operations required per tree operation to O(k) where k is the number of non-empty elements in the tree.
 
 [![Build Status](https://travis-ci.org/lazyledger/smt.svg?branch=master)](https://travis-ci.org/lazyledger/smt)
 [![Coverage Status](https://coveralls.io/repos/github/lazyledger/smt/badge.svg?branch=master)](https://coveralls.io/github/lazyledger/smt?branch=master)
 [![GoDoc](https://godoc.org/github.com/lazyledger/smt?status.svg)](https://godoc.org/github.com/lazyledger/smt)
 
-Thanks to Vitalik Buterin who provided the [original Python prototype](https://ethresear.ch/t/data-availability-proof-friendly-state-tree-transitions/1453/6) that some of this library is inspired from.
-
 ## Example
 
 ```go

bulk_test.go

@@ -0,0 +1,136 @@
+package smt
+
+import (
+	"bytes"
+	"crypto/sha256"
+	"math/rand"
+	"reflect"
+	"testing"
+)
+
+// Test all tree operations in bulk.
+func TestSparseMerkleTree(t *testing.T) {
+	for i := 0; i < 5; i++ {
+		// Test more inserts/updates than deletions.
+		bulkOperations(t, 200, 100, 100, 50)
+	}
+	for i := 0; i < 5; i++ {
+		// Test extreme deletions.
+		bulkOperations(t, 200, 100, 100, 500)
+	}
+}
+
+// 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) {
+	sm := NewSimpleMap()
+	smt := NewSparseMerkleTree(sm, sha256.New())
+
+	max := insert + update + delete
+	kv := make(map[string]string)
+
+	for i := 0; i < operations; i++ {
+		n := rand.Intn(max)
+		if n < insert { // Insert
+			keyLen := 16 + rand.Intn(32)
+			key := make([]byte, keyLen)
+			rand.Read(key)
+
+			valLen := 1 + rand.Intn(64)
+			val := make([]byte, valLen)
+			rand.Read(val)
+
+			kv[string(key)] = string(val)
+			_, err := smt.Update(key, val)
+			if err != nil {
+				t.Errorf("error: %v", err)
+			}
+		} else if n > insert && n < insert+update { // Update
+			keys := reflect.ValueOf(kv).MapKeys()
+			if len(keys) == 0 {
+				continue
+			}
+			key := []byte(keys[rand.Intn(len(keys))].Interface().(string))
+
+			valLen := 1 + rand.Intn(64)
+			val := make([]byte, valLen)
+			rand.Read(val)
+
+			kv[string(key)] = string(val)
+			_, err := smt.Update(key, val)
+			if err != nil {
+				t.Errorf("error: %v", err)
+			}
+		} else { // Delete
+			keys := reflect.ValueOf(kv).MapKeys()
+			if len(keys) == 0 {
+				continue
+			}
+			key := []byte(keys[rand.Intn(len(keys))].Interface().(string))
+
+			kv[string(key)] = ""
+			_, err := smt.Update(key, defaultValue)
+			if err != nil {
+				t.Errorf("error: %v", err)
+			}
+		}
+
+		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))
+		if err != nil {
+			t.Errorf("error: %v", err)
+		}
+		if !bytes.Equal([]byte(v), value) {
+			t.Error("got incorrect value when bulk testing operations")
+		}
+
+		// 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")
+		}
+		compactProof, err := smt.ProveCompact([]byte(k))
+		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 v == "" {
+			continue
+		}
+
+		// Check that the key is at the correct height in the tree.
+		largestCommonPrefix := 0
+		for k2, v2 := range *kv {
+			if v2 == "" {
+				continue
+			}
+			commonPrefix := countCommonPrefix(smt.th.path([]byte(k)), smt.th.path([]byte(k2)))
+			if commonPrefix != smt.depth() && commonPrefix > largestCommonPrefix {
+				largestCommonPrefix = commonPrefix
+			}
+		}
+		sideNodes, _, _, err := smt.sideNodesForRoot(smt.th.path([]byte(k)), smt.Root())
+		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")
+		}
+	}
+}

deepsubtree.go

@@ -1,31 +1,48 @@
 package smt
 
-import(
-    "hash"
+import (
+	"hash"
 )
 
+// BadProofError is returned when an invalid Merkle proof is supplied.
+type BadProofError struct{}
+
+func (e *BadProofError) Error() string {
+	return "bad proof"
+}
+
 // DeepSparseMerkleSubTree is a deep Sparse Merkle subtree for working on only a few leafs.
 type DeepSparseMerkleSubTree struct {
-    *SparseMerkleTree
+	*SparseMerkleTree
 }
 
 // NewDeepSparseMerkleSubTree creates a new deep Sparse Merkle subtree on an empty MapStore.
-func NewDeepSparseMerkleSubTree(ms MapStore, hasher hash.Hash) *DeepSparseMerkleSubTree {
-    smt := &SparseMerkleTree{
-        hasher: hasher,
-        ms: ms,
-    }
+func NewDeepSparseMerkleSubTree(ms MapStore, hasher hash.Hash, root []byte) *DeepSparseMerkleSubTree {
+	smt := &SparseMerkleTree{
+		th: *newTreeHasher(hasher),
+		ms: ms,
+	}
 
-    return &DeepSparseMerkleSubTree{SparseMerkleTree: smt}
+	smt.SetRoot(root)
+
+	return &DeepSparseMerkleSubTree{SparseMerkleTree: smt}
 }
 
-// AddBranches adds new branches to the tree.
-// These branches are generated by smt.ProveForRoot, and should be verified by VerifyProof first.
-// Set updateRoot to true if the current root of the tree should be updated.
-func (dsmst *DeepSparseMerkleSubTree) AddBranches(proof [][]byte, key []byte, value []byte, updateRoot bool) ([]byte, error) {
-    newRoot, err := dsmst.updateWithSideNodes(dsmst.digest(key), value, proof)
-    if err == nil && updateRoot {
-        dsmst.SetRoot(newRoot)
-    }
-    return newRoot, err
+// AddBranch adds a branch to the tree.
+// These branches are generated by smt.ProveForRoot.
+// If the proof is invalid, a BadProofError is returned.
+func (dsmst *DeepSparseMerkleSubTree) AddBranch(proof SparseMerkleProof, key []byte, value []byte) error {
+	result, updates := verifyProofWithUpdates(proof, dsmst.Root(), key, value, dsmst.th.hasher)
+	if !result {
+		return &BadProofError{}
+	}
+
+	for _, update := range updates {
+		err := dsmst.ms.Put(update[0], update[1])
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
 }

deepsubtree_test.go

@@ -1,62 +1,104 @@
 package smt
 
-import(
-    "crypto/sha256"
-    "testing"
-    "bytes"
+import (
+	"bytes"
+	"crypto/sha256"
+	"testing"
 )
 
-func TestDeepSparseMerkleSubTree(t *testing.T) {
-    dsmst := NewDeepSparseMerkleSubTree(NewSimpleMap(), sha256.New())
-    smt := NewSparseMerkleTree(NewSimpleMap(), sha256.New())
-
-    smt.Update([]byte("testKey1"), []byte("testValue1"))
-    smt.Update([]byte("testKey2"), []byte("testValue2"))
-    smt.Update([]byte("testKey3"), []byte("testValue3"))
-    smt.Update([]byte("testKey4"), []byte("testValue4"))
-
-    proof1, _ := smt.Prove([]byte("testKey1"))
-    proof2, _ := smt.Prove([]byte("testKey2"))
-    dsmst.AddBranches(proof1, []byte("testKey1"), []byte("testValue1"), true)
-    dsmst.AddBranches(proof2, []byte("testKey2"), []byte("testValue2"), true)
-
-    value, err := dsmst.Get([]byte("testKey1"))
-    if err != nil {
-        t.Error("returned error when getting value in deep subtree")
-    }
-    if bytes.Compare(value, []byte("testValue1")) != 0 {
-        t.Error("did not get correct value in deep subtree")
-    }
-    value, err = dsmst.Get([]byte("testKey2"))
-    if err != nil {
-        t.Error("returned error when getting value in deep subtree")
-    }
-    if bytes.Compare(value, []byte("testValue2")) != 0 {
-        t.Error("did not get correct value in deep subtree")
-    }
-
-    dsmst.Update([]byte("testKey1"), []byte("testValue3"))
-    dsmst.Update([]byte("testKey2"), []byte("testValue4"))
-
-    value, err = dsmst.Get([]byte("testKey1"))
-    if err != nil {
-        t.Error("returned error when getting value in deep subtree")
-    }
-    if bytes.Compare(value, []byte("testValue3")) != 0 {
-        t.Error("did not get correct value in deep subtree")
-    }
-    value, err = dsmst.Get([]byte("testKey2"))
-    if err != nil {
-        t.Error("returned error when getting value in deep subtree")
-    }
-    if bytes.Compare(value, []byte("testValue4")) != 0 {
-        t.Error("did not get correct value in deep subtree")
-    }
-
-    smt.Update([]byte("testKey1"), []byte("testValue3"))
-    smt.Update([]byte("testKey2"), []byte("testValue4"))
-
-    if bytes.Compare(smt.Root(), dsmst.Root()) != 0 {
-        t.Error("roots of identical standard tree and subtree do not match")
-    }
+func TestDeepSparseMerkleSubTreeBasic(t *testing.T) {
+	smt := NewSparseMerkleTree(NewSimpleMap(), sha256.New())
+
+	smt.Update([]byte("testKey1"), []byte("testValue1"))
+	smt.Update([]byte("testKey2"), []byte("testValue2"))
+	smt.Update([]byte("testKey3"), []byte("testValue3"))
+	smt.Update([]byte("testKey4"), []byte("testValue4"))
+	smt.Update([]byte("testKey6"), []byte("testValue6"))
+
+	proof1, _ := smt.Prove([]byte("testKey1"))
+	proof2, _ := smt.Prove([]byte("testKey2"))
+	proof5, _ := smt.Prove([]byte("testKey5"))
+
+	dsmst := NewDeepSparseMerkleSubTree(NewSimpleMap(), sha256.New(), smt.Root())
+	dsmst.AddBranch(proof1, []byte("testKey1"), []byte("testValue1"))
+	dsmst.AddBranch(proof2, []byte("testKey2"), []byte("testValue2"))
+	dsmst.AddBranch(proof5, []byte("testKey5"), defaultValue)
+
+	value, err := dsmst.Get([]byte("testKey1"))
+	if err != nil {
+		t.Error("returned error when getting value in deep subtree")
+	}
+	if bytes.Compare(value, []byte("testValue1")) != 0 {
+		t.Error("did not get correct value in deep subtree")
+	}
+	value, err = dsmst.Get([]byte("testKey2"))
+	if err != nil {
+		t.Error("returned error when getting value in deep subtree")
+	}
+	if bytes.Compare(value, []byte("testValue2")) != 0 {
+		t.Error("did not get correct value in deep subtree")
+	}
+	value, err = dsmst.Get([]byte("testKey5"))
+	if err != nil {
+		t.Error("returned error when getting value in deep subtree")
+	}
+	if bytes.Compare(value, defaultValue) != 0 {
+		t.Error("did not get correct value in deep subtree")
+	}
+	value, err = dsmst.Get([]byte("testKey6"))
+	if err == nil {
+		t.Error("did not error when getting non-added value in deep subtree")
+	}
+
+	dsmst.Update([]byte("testKey1"), []byte("testValue3"))
+	dsmst.Update([]byte("testKey2"), defaultValue)
+	dsmst.Update([]byte("testKey5"), []byte("testValue5"))
+
+	value, err = dsmst.Get([]byte("testKey1"))
+	if err != nil {
+		t.Error("returned error when getting value in deep subtree")
+	}
+	if bytes.Compare(value, []byte("testValue3")) != 0 {
+		t.Error("did not get correct value in deep subtree")
+	}
+	value, err = dsmst.Get([]byte("testKey2"))
+	if err != nil {
+		t.Error("returned error when getting value in deep subtree")
+	}
+	if bytes.Compare(value, defaultValue) != 0 {
+		t.Error("did not get correct value in deep subtree")
+	}
+	value, err = dsmst.Get([]byte("testKey5"))
+	if err != nil {
+		t.Error("returned error when getting value in deep subtree")
+	}
+	if bytes.Compare(value, []byte("testValue5")) != 0 {
+		t.Error("did not get correct value in deep subtree")
+	}
+
+	smt.Update([]byte("testKey1"), []byte("testValue3"))
+	smt.Update([]byte("testKey2"), defaultValue)
+	smt.Update([]byte("testKey5"), []byte("testValue5"))
+
+	if bytes.Compare(smt.Root(), dsmst.Root()) != 0 {
+		t.Error("roots of identical standard tree and subtree do not match")
+	}
+}
+
+func TestDeepSparseMerkleSubTreeBadInput(t *testing.T) {
+	smt := NewSparseMerkleTree(NewSimpleMap(), sha256.New())
+
+	smt.Update([]byte("testKey1"), []byte("testValue1"))
+	smt.Update([]byte("testKey2"), []byte("testValue2"))
+	smt.Update([]byte("testKey3"), []byte("testValue3"))
+	smt.Update([]byte("testKey4"), []byte("testValue4"))
+
+	badProof, _ := smt.Prove([]byte("testKey1"))
+	badProof.SideNodes[0][0] = byte(0)
+
+	dsmst := NewDeepSparseMerkleSubTree(NewSimpleMap(), sha256.New(), smt.Root())
+	err := dsmst.AddBranch(badProof, []byte("testKey1"), []byte("testValue1"))
+	if _, ok := err.(*BadProofError); !ok {
+		t.Error("did not return BadProofError for bad proof input")
+	}
 }