[Persistence] Change root hash calculation to use an SMT

persistence/docs/PROTOCOL_STATE_HASH.md

@@ -23,14 +23,14 @@ This document defines how Pocket V1 takes a snapshot of its world state. An intr
 
 ### Infrastructural Components
 
-| Component             | Data Type                             | Implementation Options - Examples                      | Implementation Selected - Current | Example             | Use Case                                                                         |
-| --------------------- | ------------------------------------- | ------------------------------------------------------ | --------------------------------- | ------------------- | -------------------------------------------------------------------------------- |
-| Data Tables           | SQL Database / Engine                 | MySQL, SQLite, PostgreSQL                              | PostgresSQL                       | Validator SQL Table | Validating & updating information when applying a transaction                    |
-| Merkle Trees          | Merkle Trie backed by Key-Value Store | Celestia's SMT, Libra's JMT, Cosmos' IAVL, Verkle Tree | Celestia's SMT                    | Fisherman Trie      | Maintains the state of all account based trees                                   |
-| Blocks                | Serialization Codec                   | Amino, Protobuf, Thrift, Avro                          | Protobuf                          | Block protobuf      | Serialized and inserted into the Block Store                                     |
-| Objects (e.g. Actors) | Serialization Codec                   | Amino, Protobuf, Thrift, Avro                          | Protobuf                          | Servicer protobuf   | Serialized and inserted into the corresponding Tree                              |
-| Block Store           | Key Value Store                       | LevelDB, BadgerDB, RocksDB, BoltDB                     | BadgerDb                          | Block Store         | Maintains a key-value store of the blockchain blocks                             |
-| Transaction Indexer   | Key Value Store                       | LevelDB, BadgerDB, RocksDB, BoltDB                     | BadgerDB                          | Tx Indexer          | Indexes transactions in different ways for fast queries, presence checks, etc... |
+| Component             | Data Type                             | Implementation Options - Examples                      | Implementation Selected - Current   | Example             | Use Case                                                                         |
+| --------------------- | ------------------------------------- | ------------------------------------------------------ | ----------------------------------- | ------------------- | -------------------------------------------------------------------------------- |
+| Data Tables           | SQL Database / Engine                 | MySQL, SQLite, PostgreSQL                              | PostgresSQL                         | Validator SQL Table | Validating & updating information when applying a transaction                    |
+| Merkle Trees          | Merkle Trie backed by Key-Value Store | Celestia's SMT, Libra's JMT, Cosmos' IAVL, Verkle Tree | Pocket's SMT (Forked from Celestia) | Fisherman Trie      | Maintains the state of all account based trees                                   |
+| Blocks                | Serialization Codec                   | Amino, Protobuf, Thrift, Avro                          | Protobuf                            | Block protobuf      | Serialized and inserted into the Block Store                                     |
+| Objects (e.g. Actors) | Serialization Codec                   | Amino, Protobuf, Thrift, Avro                          | Protobuf                            | Servicer protobuf   | Serialized and inserted into the corresponding Tree                              |
+| Block Store           | Key Value Store                       | LevelDB, BadgerDB, RocksDB, BoltDB                     | BadgerDb                            | Block Store         | Maintains a key-value store of the blockchain blocks                             |
+| Transaction Indexer   | Key Value Store                       | LevelDB, BadgerDB, RocksDB, BoltDB                     | BadgerDB                            | Tx Indexer          | Indexes transactions in different ways for fast queries, presence checks, etc... |
 
 ### Block Proto
 
@@ -40,8 +40,6 @@ The block protobuf that is serialized and store in the block store can be found
 
 An individual Merkle Tree is created for each type of actor, record or data type. Each of these is backed by its own key-value store.
 
-Note that the order in which the trees are defined (found in `persistence/state.go`) is important since it determines how the state hash is computed. _TODO(#361): Consider specifying the oder in a `.proto` `enum` rather than a `.go` `iota`._
-
 **Actor Merkle Trees**:
 
 - Applications
@@ -71,8 +69,8 @@ This flow shows the interaction between the `PostgresDB` and `MerkleTrees` liste
 3. Serialize each record using the corresponding underlying protobuf
 4. Insert the serialized record into the corresponding tree (which is back by a key-value store)
 5. Compute the root hash of each tree
-6. Aggregate all the root hashes by concatenating them together
-7. Compute the new `stateHash` by taking a hash of the concatenated hash list
+6. Insert the name of the tree and its root hash into the root tree
+7. Compute the new `stateHash` by hex encoding the root tree's root hash
 
 ```mermaid
 sequenceDiagram
@@ -89,15 +87,15 @@ sequenceDiagram
         end
         P->>+PKV: GetRoot()
         PKV->>-P: rootHash
+        P->>P: rootTree.Update(stateTreeName, rootHash)
     end
 
-    P->>P: stateHash = hash(concat(rootHashes))
+    P->>P: stateHash = hex(rootTree.Root())
+
     activate P
     deactivate P
 ```
 
-_IMPORTANT: The order in which the `rootHashes` are concatenated is based on the definition in which the trees are ordered in within `state.go`._
-
 ## Store Block (Commit)
 
 When the `Commit(proposer, quorumCert)` function is invoked, the current context is committed to disk. The `PersistenceContext` does the following:

persistence/module.go

@@ -235,6 +235,10 @@ func (m *persistenceModule) GetTxIndexer() indexer.TxIndexer {
 	return m.txIndexer
 }
 
+func (m *persistenceModule) GetTreeStore() modules.TreeStoreModule {
+	return m.stateTrees
+}
+
 func (m *persistenceModule) GetNetworkID() string {
 	return m.networkId
 }

persistence/test/state_test.go

@@ -42,9 +42,9 @@ func TestStateHash_DeterministicStateWhenUpdatingAppStake(t *testing.T) {
 	// logic changes, these hashes will need to be updated based on the test output.
 	// TODO: Add an explicit updateSnapshots flag to the test to make this more clear.
 	stateHashes := []string{
-		"6a30f096c86793de894388aad171e84c5ce766cbd82c5a5d36ca53d6c99ac041",
-		"c5840401da7028948f6d025867249fb9f9a9e738b36158669a64746e5b4f3ed2",
-		"209a2bda7e9a4495f45281e726777b33f54ac61868c8d8a059719cb5cefd2f75",
+		"fb1c1b2da242eb6148884e1f11c184c673963b6fcb59feea4ea51c9840e4b56c",
+		"9ae40f9fd0864c01760c19d6688d1c7bba5ad746fc3f7123cf071b142c2a302a",
+		"5c13743e9d29e83cbff1301e3750c7583e365b2c3fd81b643ed8db985f46268f",
 	}
 
 	stakeAmount := initialStakeAmount

persistence/trees/module.go

@@ -44,30 +44,47 @@ func (t *treeStore) setupTrees() error {
 		return t.setupInMemory()
 	}
 
-	t.merkleTrees = make(map[merkleTree]*smt.SMT, int(numMerkleTrees))
-	t.nodeStores = make(map[merkleTree]kvstore.KVStore, int(numMerkleTrees))
+	nodeStore, err := kvstore.NewKVStore(fmt.Sprintf("%s/%s_nodes", t.treeStoreDir, RootTreeName))
+	if err != nil {
+		return err
+	}
+	t.rootTree = &stateTree{
+		name:      RootTreeName,
+		tree:      smt.NewSparseMerkleTree(nodeStore, smtTreeHasher),
+		nodeStore: nodeStore,
+	}
+	t.merkleTrees = make(map[string]*stateTree, len(stateTreeNames))
 
-	for tree := merkleTree(0); tree < numMerkleTrees; tree++ {
-		nodeStore, err := kvstore.NewKVStore(fmt.Sprintf("%s/%s_nodes", t.treeStoreDir, merkleTreeToString[tree]))
+	for i := 0; i < len(stateTreeNames); i++ {
+		nodeStore, err := kvstore.NewKVStore(fmt.Sprintf("%s/%s_nodes", t.treeStoreDir, stateTreeNames[i]))
 		if err != nil {
 			return err
 		}
-		t.nodeStores[tree] = nodeStore
-		t.merkleTrees[tree] = smt.NewSparseMerkleTree(nodeStore, smtTreeHasher)
+		t.merkleTrees[stateTreeNames[i]] = &stateTree{
+			name:      stateTreeNames[i],
+			tree:      smt.NewSparseMerkleTree(nodeStore, smtTreeHasher),
+			nodeStore: nodeStore,
+		}
 	}
 
 	return nil
 }
 
 func (t *treeStore) setupInMemory() error {
-	t.merkleTrees = make(map[merkleTree]*smt.SMT, int(numMerkleTrees))
-	t.nodeStores = make(map[merkleTree]kvstore.KVStore, int(numMerkleTrees))
-
-	for tree := merkleTree(0); tree < numMerkleTrees; tree++ {
+	nodeStore := kvstore.NewMemKVStore()
+	t.rootTree = &stateTree{
+		name:      RootTreeName,
+		tree:      smt.NewSparseMerkleTree(nodeStore, smtTreeHasher),
+		nodeStore: nodeStore,
+	}
+	t.merkleTrees = make(map[string]*stateTree, len(stateTreeNames))
+	for i := 0; i < len(stateTreeNames); i++ {
 		nodeStore := kvstore.NewMemKVStore() // For testing, `smt.NewSimpleMap()` can be used as well
-		t.nodeStores[tree] = nodeStore
-		t.merkleTrees[tree] = smt.NewSparseMerkleTree(nodeStore, smtTreeHasher)
+		t.merkleTrees[stateTreeNames[i]] = &stateTree{
+			name:      stateTreeNames[i],
+			tree:      smt.NewSparseMerkleTree(nodeStore, smtTreeHasher),
+			nodeStore: nodeStore,
+		}
 	}
-
 	return nil
 }