[Enhancement] Add MerkleRoot type alias for []byte with Sum method (#37)

docs/merkle-sum-trie.md

@@ -10,6 +10,7 @@
     + [General Trie Structure](#general-trie-structure)
     + [Binary Sum Digests](#binary-sum-digests)
 - [Sum](#sum)
+- [Roots](#roots)
 - [Nil Values](#nil-values)
 
 <!-- tocstop -->
@@ -248,6 +249,15 @@ graph TB
 The `Sum()` function adds functionality to easily retrieve the trie's current
 sum as a `uint64`.
 
+## Roots
+
+The root of the tree is a slice of bytes. `MerkleRoot` is an alias for `[]byte`.
+This design enables easily passing around the data (e.g. on-chain)
+while maintaining primitive usage in different use cases (e.g. proofs).
+
+`MerkleRoot` provides helpers, such as retrieving the `Sum(sumTrie bool)uint64`
+to interface with data it captures.
+
 ## Nil Values
 
 A `nil` value and `0` weight is the same as the placeholder value and default

docs/smt.md

@@ -17,6 +17,7 @@
 - [Values](#values)
   * [Nil values](#nil-values)
 - [Hashers & Digests](#hashers--digests)
+- [Roots](#roots)
 - [Proofs](#proofs)
   * [Verification](#verification)
   * [Closest Proof](#closest-proof)
@@ -26,9 +27,6 @@
 - [Database](#database)
   * [Database Submodules](#database-submodules)
     + [SimpleMap](#simplemap)
-    + [Badger](#badger)
-  * [Data Loss](#data-loss)
-- [Sparse Merkle Sum Trie](#sparse-merkle-sum-trie)
 
 <!-- tocstop -->
 
@@ -342,6 +340,16 @@ graph TD
 	VH --ValueHash-->L
 ```
 
+## Roots
+
+The root of the tree is a slice of bytes. `MerkleRoot` is an alias for `[]byte`.
+This design enables easily passing around the data (e.g. on-chain)
+while maintaining primitive usage in different use cases (e.g. proofs).
+
+`MerkleRoot` provides helpers, such as retrieving the `Sum(sumTrie bool)uint64`
+to interface with data it captures. However, for the SMT it **always** panics,
+as there is no sum.
+
 ## Proofs
 
 The `SparseMerkleProof` type contains the information required for inclusion
@@ -469,33 +477,4 @@ the [`kvstore`](../kvstore/) directory.
 
 #### SimpleMap
 
-This library defines the `SimpleMap` interface which is implemented as an
-extremely simple in-memory key-value store.
-
-Although it is a submodule, it is ideal for simple, testing or non-production
-use cases. It is used in the tests throughout the library.
-
-See [simplemap.go](../kvstore/simplemap/simplemap.go) for the implementation
-details.
-
-#### Badger
-
-This library defines the `BadgerStore` interface which is implemented as a
-wrapper around the [BadgerDB](https://github.com/dgraph-io/badger) v4 key-value
-database. It's interface exposes numerous extra methods not used by the trie,
-However it can still be used as a node-store with both in-memory and persistent
-options.
-
-See [badger-store.md](./badger-store.md.md) for the details of the implementation.
-
-### Data Loss
-
-In the event of a system crash or unexpected failure of the program utilising
-the SMT, if the `Commit()` function has not been called, any changes to the trie
-will be lost. This is due to the underlying database not being changed **until**
-the `Commit()` function is called and changes are persisted.
-
-## Sparse Merkle Sum Trie
-
-This library also implements a Sparse Merkle Sum Trie (SMST), the documentation
-for which can be found [here](./merkle-sum-trie.md).
+This l

root_test.go

@@ -0,0 +1,82 @@
+package smt_test
+
+import (
+	"crypto/sha256"
+	"crypto/sha512"
+	"encoding/binary"
+	"fmt"
+	"hash"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/pokt-network/smt"
+	"github.com/pokt-network/smt/kvstore/simplemap"
+)
+
+func TestMerkleRoot_TrieTypes(t *testing.T) {
+	tests := []struct {
+		desc          string
+		sumTree       bool
+		hasher        hash.Hash
+		expectedPanic string
+	}{
+		{
+			desc:          "successfully: gets sum of sha256 hasher SMST",
+			sumTree:       true,
+			hasher:        sha256.New(),
+			expectedPanic: "",
+		},
+		{
+			desc:          "successfully: gets sum of sha512 hasher SMST",
+			sumTree:       true,
+			hasher:        sha512.New(),
+			expectedPanic: "",
+		},
+		{
+			desc:          "failure: panics for sha256 hasher SMT",
+			sumTree:       false,
+			hasher:        sha256.New(),
+			expectedPanic: "roo#sum: not a merkle sum trie",
+		},
+		{
+			desc:          "failure: panics for sha512 hasher SMT",
+			sumTree:       false,
+			hasher:        sha512.New(),
+			expectedPanic: "roo#sum: not a merkle sum trie",
+		},
+	}
+
+	nodeStore := simplemap.NewSimpleMap()
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.desc, func(t *testing.T) {
+			t.Cleanup(func() {
+				require.NoError(t, nodeStore.ClearAll())
+			})
+			if tt.sumTree {
+				trie := smt.NewSparseMerkleSumTrie(nodeStore, tt.hasher)
+				for i := uint64(0); i < 10; i++ {
+					require.NoError(t, trie.Update([]byte(fmt.Sprintf("key%d", i)), []byte(fmt.Sprintf("value%d", i)), i))
+				}
+				root := trie.Root()
+				require.Equal(t, root.Sum(), getSumBzHelper(t, root))
+				return
+			}
+			trie := smt.NewSparseMerkleTrie(nodeStore, tt.hasher)
+			for i := 0; i < 10; i++ {
+				require.NoError(t, trie.Update([]byte(fmt.Sprintf("key%d", i)), []byte(fmt.Sprintf("value%d", i))))
+			}
+			if panicStr := recover(); panicStr != nil {
+				require.Equal(t, tt.expectedPanic, panicStr)
+			}
+		})
+	}
+}
+
+func getSumBzHelper(t *testing.T, r []byte) uint64 {
+	sumSize := len(r) % 32
+	sumBz := make([]byte, sumSize)
+	copy(sumBz[:], []byte(r)[len([]byte(r))-sumSize:])
+	return binary.BigEndian.Uint64(sumBz[:])
+}

smst.go

@@ -112,14 +112,12 @@ func (smst *SMST) Commit() error {
 }
 
 // Root returns the root hash of the trie with the total sum bytes appended
-func (smst *SMST) Root() []byte {
+func (smst *SMST) Root() MerkleRoot {
 	return smst.SMT.Root() // [digest]+[binary sum]
 }
 
 // Sum returns the uint64 sum of the entire trie
 func (smst *SMST) Sum() uint64 {
-	var sumBz [sumSize]byte
 	digest := smst.Root()
-	copy(sumBz[:], digest[len(digest)-sumSize:])
-	return binary.BigEndian.Uint64(sumBz[:])
+	return digest.Sum()
 }

smt.go

@@ -675,7 +675,7 @@ func (smt *SMT) commit(node trieNode) error {
 }
 
 // Root returns the root hash of the trie
-func (smt *SMT) Root() []byte {
+func (smt *SMT) Root() MerkleRoot {
 	return hashNode(smt.Spec(), smt.trie)
 }
 

types.go

@@ -1,6 +1,7 @@
 package smt
 
 import (
+	"encoding/binary"
 	"hash"
 )
 
@@ -14,6 +15,21 @@ var (
 	defaultSum   [sumSize]byte
 )
 
+// MerkleRoot is a type alias for a byte slice returned from the Root method
+type MerkleRoot []byte
+
+// Sum returns the uint64 sum of the merkle root, it checks the length of the
+// merkle root and if it is no the same as the size of the SMST's expected
+// root hash it will panic.
+func (r MerkleRoot) Sum() uint64 {
+	if len(r)%32 == 0 {
+		panic("roo#sum: not a merkle sum trie")
+	}
+	var sumbz [sumSize]byte
+	copy(sumbz[:], []byte(r)[len([]byte(r))-sumSize:])
+	return binary.BigEndian.Uint64(sumbz[:])
+}
+
 // SparseMerkleTrie represents a Sparse Merkle Trie.
 type SparseMerkleTrie interface {
 	// Update inserts a value into the SMT.
@@ -23,7 +39,7 @@ type SparseMerkleTrie interface {
 	// Get descends the trie to access a value. Returns nil if key is not present.
 	Get(key []byte) ([]byte, error)
 	// Root computes the Merkle root digest.
-	Root() []byte
+	Root() MerkleRoot
 	// Prove computes a Merkle proof of inclusion or exclusion of a key.
 	Prove(key []byte) (*SparseMerkleProof, error)
 	// ProveClosest computes a Merkle proof of inclusion for a key in the trie
@@ -46,7 +62,7 @@ type SparseMerkleSumTrie interface {
 	// Get descends the trie to access a value. Returns nil if key is not present.
 	Get(key []byte) ([]byte, uint64, error)
 	// Root computes the Merkle root digest.
-	Root() []byte
+	Root() MerkleRoot
 	// Sum computes the total sum of the Merkle trie
 	Sum() uint64
 	// Prove computes a Merkle proof of inclusion or exclusion of a key.