[Exploration] Using Celestia's SMT for Pocket Network V1?

.gitignore

@@ -0,0 +1 @@
+.vscode

Makefile

@@ -0,0 +1,33 @@
+.SILENT:
+
+help:
+	printf "Available targets\n\n"
+	awk '/^[a-zA-Z\-\_0-9]+:/ { \
+		helpMessage = match(lastLine, /^## (.*)/); \
+		if (helpMessage) { \
+			helpCommand = substr($$1, 0, index($$1, ":")-1); \
+			helpMessage = substr(lastLine, RSTART + 3, RLENGTH); \
+			printf "%-30s %s\n", helpCommand, helpMessage; \
+		} \
+	} \
+	{ lastLine = $$0 }' $(MAKEFILE_LIST)
+
+.PHONY: test_all
+## Run all the unit tests
+test_all:
+	go test -v -count=1 ./...
+
+.PHONY: test_smt
+## Run all the ^TestSparseMerkleTree unit tests
+test_smt:
+	go test -v -count=1 -run TestSparseMerkleTree ./...
+
+.PHONY: test_th
+## Run all the ^TestTreeHasher unit tests
+test_th:
+	go test -v -count=1 -run TestTreeHasher ./...
+
+.PHONY: test_ms
+## Run all the ^TestMapStore unit tests
+test_ms:
+	go test -v -count=1 -run TestMapStore ./...

README.md

@@ -1,44 +1,93 @@
-# smt
+# Sparse Merkle Tree (smt)
 
-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][libra whitepaper], 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.
+A Go library that implements a Sparse Merkle Tree for a key-value map. The tree implements the same optimisations specified in the [Jellyfish Merkle Tree whitepaper][jmt whitepaper] originally designed for the [Libra blockchain][libra whitepaper]. It reduces the number of hash operations required per tree operation to `O(k)` where `k` is the number of non-empty elements in the tree.
 
 [![Tests](https://github.com/celestiaorg/smt/actions/workflows/test.yml/badge.svg)](https://github.com/celestiaorg/smt/actions/workflows/test.yml)
 [![codecov](https://codecov.io/gh/celestiaorg/smt/branch/master/graph/badge.svg?token=U3GGEDSA94)](https://codecov.io/gh/celestiaorg/smt)
 [![GoDoc](https://godoc.org/github.com/celestiaorg/smt?status.svg)](https://godoc.org/github.com/celestiaorg/smt)
 
+## Installation
+
+```bash
+go get github.com/celestiaorg/smt@master
+```
+
 ## Example
 
 ```go
 package main
 
 import (
-	"crypto/sha256"
-	"fmt"
+    "crypto/sha256"
+    "fmt"
 
-	"github.com/celestiaorg/smt"
+    "github.com/celestiaorg/smt"
 )
 
 func main() {
-	// Initialise two new key-value store to store the nodes and values of the tree
-	nodeStore := smt.NewSimpleMap()
-	valueStore := smt.NewSimpleMap()
-	// Initialise the tree
-	tree := smt.NewSparseMerkleTree(nodeStore, valueStore, sha256.New())
-
-	// Update the key "foo" with the value "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()) {
-		fmt.Println("Proof verification succeeded.")
-	} else {
-		fmt.Println("Proof verification failed.")
-	}
+    // Initialise 2 new key-value store to stores the nodes and values of the tree
+    nodeStore := smt.NewSimpleMap() // Mapping from hash -> data;
+    valueStore := smt.NewSimpleMap() // Mapping from node_path -> node_value; a path can be retrieved using the digest of the key
+
+    // Initialise the smt
+    tree := smt.NewSparseMerkleTree(nodeStore, valueStore, sha256.New())
+
+    // Update the key "foo" with the value "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()) {
+        fmt.Println("Proof verification succeeded.")
+    } else {
+        fmt.Println("Proof verification failed.")
+    }
 }
 ```
 
+## Development
+
+Run `make` to see all the options available
+
+## General Improvements / TODOs
+
+- [ ] Use the `require` test module to simplify unit tests; can be done with a single clever regex find+replace
+- [ ] Create types for `sideNodes`, `root`, etc...
+- [ ] Add an interface for `SparseMerkleProof` so we can return nils and not access vars directly
+- [ ] Add an interface for `SparseMerkleTree` so it's clear how we should interact with it
+- [ ] If we create an interface for `TreeHasher`, we can embed it in `SparseMerkleTree` and then avoid the need to write things like `smt.th.path(...)` everywhere and use `smt.path(...)` directly.
+- [ ] Consider splitting `smt.go` into `smt_ops.go` and `smt_proofs.go`
+- [ ] Functions like `sideNodesForRoot` and `updateWithSideNodes` need to be split into smaller more compartmentalized functions
+
 [libra whitepaper]: https://diem-developers-components.netlify.app/papers/the-diem-blockchain/2020-05-26.pdf
+[jmt whitepaper]: https://developers.diem.com/papers/jellyfish-merkle-tree/2021-01-14.pdf
+
+### [Delete me later] personal checklist
+
+- [x] ├── LICENSE
+- [x] ├── Makefile
+- [x] ├── README.md
+- [ ] ├── bench_test.go
+- [ ] ├── bulk_test.go
+- [ ] ├── deepsubtree.go
+- [ ] ├── deepsubtree_test.go
+- [ ] ├── fuzz
+- [ ] │   ├── delete
+- [ ] │   │   └── fuzz.go
+- [ ] │   └── fuzz.go
+- [x] ├── go.mod
+- [x] ├── go.sum
+- [x] ├── mapstore.go
+- [x] ├── mapstore_test.go
+- [x] ├── options.go
+- [ ] ├── oss-fuzz-build.sh
+- [ ] ├── proofs.go
+- [ ] ├── proofs_test.go
+- [x] ├── smt.go
+- [ ] ├── smt_test.go
+- [x] ├── treehasher.go
+- [x] ├── treehasher_test.go
+- [x] └── utils.go

bulk_test.go

@@ -6,6 +6,8 @@ import (
 	"math/rand"
 	"reflect"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 // Test all tree operations in bulk.
@@ -21,7 +23,7 @@ func TestSparseMerkleTree(t *testing.T) {
 }
 
 // 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) {
+func bulkOperations(t *testing.T, operations, insert, update, delete int) (*SparseMerkleTree, *SimpleMap, *SimpleMap, map[string]string) {
 	smn, smv := NewSimpleMap(), NewSimpleMap()
 	smt := NewSparseMerkleTree(smn, smv, sha256.New())
 
@@ -74,12 +76,13 @@ func bulkOperations(t *testing.T, operations int, insert int, update int, delete
 			}
 		}
 
-		bulkCheckAll(t, smt, &kv)
+		bulkCheckAll(t, smt, kv)
 	}
+	return smt, smn, smv, kv
 }
 
-func bulkCheckAll(t *testing.T, smt *SparseMerkleTree, kv *map[string]string) {
-	for k, v := range *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)
@@ -109,28 +112,36 @@ func bulkCheckAll(t *testing.T, smt *SparseMerkleTree, kv *map[string]string) {
 		}
 
 		// 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
-			}
+		largestCommonPrefix := getLargestCommonPrefix(t, smt, kv, k)
+		numSideNodes := getNumSideNodes(t, smt, kv, k)
+		if (numSideNodes != largestCommonPrefix+1) && numSideNodes != 0 && largestCommonPrefix != 0 {
+			t.Error("leaf is at unexpected height")
 		}
-		sideNodes, _, _, _, err := smt.sideNodesForRoot(smt.th.path([]byte(k)), smt.Root(), false)
-		if err != nil {
-			t.Errorf("error: %v", err)
+	}
+}
+
+func getNumSideNodes(t *testing.T, smt *SparseMerkleTree, kv map[string]string, key string) (numSideNodes int) {
+	path := smt.th.path([]byte(key))
+	sideNodes, _, _, _, err := smt.sideNodesForRoot(path, smt.Root())
+	require.NoError(t, err)
+	for _, v := range sideNodes {
+		if v != nil {
+			numSideNodes++
 		}
-		numSideNodes := 0
-		for _, v := range sideNodes {
-			if v != nil {
-				numSideNodes++
-			}
+	}
+	return
+}
+
+func getLargestCommonPrefix(_ *testing.T, smt *SparseMerkleTree, kv map[string]string, key string) (largestCommonPrefix int) {
+	path := smt.th.path([]byte(key))
+	for k, v := range kv {
+		if v == "" {
+			continue
 		}
-		if numSideNodes != largestCommonPrefix+1 && (numSideNodes != 0 && largestCommonPrefix != 0) {
-			t.Error("leaf is at unexpected height")
+		commonPrefix := countCommonPrefix(path, smt.th.path([]byte(k)))
+		if commonPrefix != smt.depth() && commonPrefix > largestCommonPrefix {
+			largestCommonPrefix = commonPrefix
 		}
 	}
+	return
 }

go.mod

@@ -1,3 +1,5 @@
 module github.com/celestiaorg/smt
 
 go 1.14
+
+require github.com/stretchr/testify v1.8.0

go.sum

@@ -0,0 +1,14 @@
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
+github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk=
+github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

mapstore.go

@@ -6,9 +6,9 @@ import (
 
 // MapStore is a key-value store.
 type MapStore interface {
-	Get(key []byte) ([]byte, error)     // Get gets the value for a key.
-	Set(key []byte, value []byte) error // Set updates the value for a key.
-	Delete(key []byte) error            // Delete deletes a key.
+	Get(key []byte) ([]byte, error) // Get gets the value for a key.
+	Set(key, value []byte) error    // Set updates the value for a key.
+	Delete(key []byte) error        // Delete deletes a key.
 }
 
 // InvalidKeyError is thrown when a key that does not exist is being accessed.
@@ -48,8 +48,7 @@ func (sm *SimpleMap) Set(key []byte, value []byte) error {
 
 // Delete deletes a key.
 func (sm *SimpleMap) Delete(key []byte) error {
-	_, ok := sm.m[string(key)]
-	if ok {
+	if _, ok := sm.m[string(key)]; ok {
 		delete(sm.m, string(key))
 		return nil
 	}

mapstore_test.go

@@ -6,26 +6,28 @@ import (
 	"testing"
 )
 
-func TestSimpleMap(t *testing.T) {
+func TestMapStoreSimpleMap(t *testing.T) {
 	sm := NewSimpleMap()
 	h := sha256.New()
+
 	var value []byte
 	var err error
 
 	h.Write([]byte("test"))
+	key := h.Sum(nil)
 
 	// Tests for Get.
-	_, err = sm.Get(h.Sum(nil))
+	_, err = sm.Get(key)
 	if err == nil {
 		t.Error("did not return an error when getting a non-existent key")
 	}
 
 	// Tests for Put.
-	err = sm.Set(h.Sum(nil), []byte("hello"))
+	err = sm.Set(key, []byte("hello"))
 	if err != nil {
 		t.Error("updating a key returned an error")
 	}
-	value, err = sm.Get(h.Sum(nil))
+	value, err = sm.Get(key)
 	if err != nil {
 		t.Error("getting a key returned an error")
 	}
@@ -34,11 +36,11 @@ func TestSimpleMap(t *testing.T) {
 	}
 
 	// Tests for Del.
-	err = sm.Delete(h.Sum(nil))
+	err = sm.Delete(key)
 	if err != nil {
 		t.Error("deleting a key returned an error")
 	}
-	_, err = sm.Get(h.Sum(nil))
+	_, err = sm.Get(key)
 	if err == nil {
 		t.Error("failed to delete key")
 	}