[Enhancement] Add MerkleRoot type alias for []byte with Sum method

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
+
+Roots are `[]byte` types aliases by the `MerkleRoot` type. This type has one
+method `Sum(sumTrie bool) uint64`. For the SMST this method is used by its own
+`Sum()` method to return the total sum of the trie.
+
+The `MerkleRoot` type being an alias means it can be used in place of the
+`[]byte` type. Specifically for proofs.
+
 ## 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)
@@ -342,6 +343,14 @@ graph TD
 	VH --ValueHash-->L
 ```
 
+## Roots
+
+Roots are `[]byte` types aliases by the `MerkleRoot` type. This type has one
+method `Sum(sumTrie bool) uint64`. For the SMT this method **always** returns 0.
+
+The `MerkleRoot` type being an alias means it can be used in place of the
+`[]byte` type. Specifically for proofs.
+
 ## Proofs
 
 The `SparseMerkleProof` type contains the information required for inclusion

root_test.go

@@ -0,0 +1,39 @@
+package smt_test
+
+import (
+	"crypto/sha256"
+	"encoding/binary"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/pokt-network/smt"
+	"github.com/pokt-network/smt/kvstore/simplemap"
+)
+
+func TestMerkleRoot_SumTrie(t *testing.T) {
+	nodeStore := simplemap.NewSimpleMap()
+	trie := smt.NewSparseMerkleSumTrie(nodeStore, sha256.New())
+	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(true), getSumBzHelper(t, root))
+}
+
+func TestMerkleRoot_Trie(t *testing.T) {
+	nodeStore := simplemap.NewSimpleMap()
+	trie := smt.NewSparseMerkleTrie(nodeStore, sha256.New())
+	for i := 0; i < 10; i++ {
+		require.NoError(t, trie.Update([]byte(fmt.Sprintf("key%d", i)), []byte(fmt.Sprintf("value%d", i))))
+	}
+	root := trie.Root()
+	require.Equal(t, root.Sum(false), uint64(0))
+}
+
+func getSumBzHelper(t *testing.T, r []byte) uint64 {
+	var sumbz [8]byte                            // Using sha256
+	copy(sumbz[:], []byte(r)[len([]byte(r))-8:]) // Using sha256 so - 8 bytes
+	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(true)
 }

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,20 @@ 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, if sumTrie is true
+// otherwise it returns 0, as it would result in undefined behaviour.
+func (r MerkleRoot) Sum(sumTrie bool) uint64 {
+	if sumTrie {
+		var sumbz [sumSize]byte
+		copy(sumbz[:], []byte(r)[len([]byte(r))-sumSize:])
+		return binary.BigEndian.Uint64(sumbz[:])
+	}
+	return 0
+}
+
 // SparseMerkleTrie represents a Sparse Merkle Trie.
 type SparseMerkleTrie interface {
 	// Update inserts a value into the SMT.
@@ -23,7 +38,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 +61,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.