Split Storage::validate_and_commit into two separate methods `Stora…

…ge::calculate_root_hash...` and `Storage::commit` (#752)

* Extend storage interface

* fix

* add calculate_state_root_and_node_batch

* Add comment

* code review feedback

* Add StateUpdate to Storage

module-system/sov-modules-stf-template/src/lib.rs

@@ -82,7 +82,6 @@ where
         witness: <Self as StateTransitionFunction<Vm, DA::BlobTransaction>>::Witness,
     ) {
         let state_checkpoint = StateCheckpoint::with_witness(self.current_storage.clone(), witness);
-
         let mut working_set = state_checkpoint.to_revertable();
 
         self.runtime.begin_slot_hook(slot_data, &mut working_set);
@@ -93,14 +92,15 @@ where
     #[cfg_attr(all(target_os = "zkvm", feature = "bench"), cycle_tracker)]
     fn end_slot(&mut self) -> (jmt::RootHash, <<C as Spec>::Storage as Storage>::Witness) {
         let (cache_log, witness) = self.checkpoint.take().unwrap().freeze();
-        let root_hash = self
+        let (root_hash, authenticated_node_batch) = self
             .current_storage
-            .validate_and_commit(cache_log, &witness)
+            .compute_state_update(cache_log, &witness)
             .expect("jellyfish merkle tree update must succeed");
 
         let mut working_set = WorkingSet::new(self.current_storage.clone());
         self.runtime.end_slot_hook(root_hash, &mut working_set);
 
+        self.current_storage.commit(&authenticated_node_batch);
         (jmt::RootHash(root_hash), witness)
     }
 }
@@ -132,11 +132,12 @@ where
             .expect("module initialization must succeed");
 
         let (log, witness) = working_set.checkpoint().freeze();
-        let genesis_hash = self
+        let (genesis_hash, node_batch) = self
             .current_storage
-            .validate_and_commit(log, &witness)
+            .compute_state_update(log, &witness)
             .expect("Storage update must succeed");
 
+        self.current_storage.commit(&node_batch);
         jmt::RootHash(genesis_hash)
     }
 

module-system/sov-state/src/prover_storage.rs

@@ -3,7 +3,7 @@ use std::marker::PhantomData;
 use std::path::Path;
 use std::sync::Arc;
 
-use jmt::storage::TreeWriter;
+use jmt::storage::{NodeBatch, TreeWriter};
 use jmt::{JellyfishMerkleTree, KeyHash, RootHash, Version};
 use sov_db::state_db::StateDB;
 
@@ -63,6 +63,7 @@ impl<S: MerkleProofSpec> Storage for ProverStorage<S> {
     type Witness = S::Witness;
     type RuntimeConfig = Config;
     type Proof = jmt::proof::SparseMerkleProof<S::Hasher>;
+    type StateUpdate = NodeBatch;
 
     fn with_config(config: Self::RuntimeConfig) -> Result<Self, anyhow::Error> {
         Self::with_path(config.path.as_path())
@@ -79,11 +80,11 @@ impl<S: MerkleProofSpec> Storage for ProverStorage<S> {
             .map(|root| root.0)
     }
 
-    fn validate_and_commit(
+    fn compute_state_update(
         &self,
         state_accesses: OrderedReadsAndWrites,
         witness: &Self::Witness,
-    ) -> Result<[u8; 32], anyhow::Error> {
+    ) -> Result<([u8; 32], Self::StateUpdate), anyhow::Error> {
         let latest_version = self.db.get_next_version() - 1;
         witness.add_hint(latest_version);
 
@@ -139,11 +140,14 @@ impl<S: MerkleProofSpec> Storage for ProverStorage<S> {
             .put_value_set(batch, next_version)
             .expect("JMT update must succeed");
 
+        Ok((new_root.0, tree_update.node_batch))
+    }
+
+    fn commit(&self, node_batch: &Self::StateUpdate) {
         self.db
-            .write_node_batch(&tree_update.node_batch)
+            .write_node_batch(node_batch)
             .expect("db write must succeed");
         self.db.inc_next_version();
-        Ok(new_root.0)
     }
 
     // Based on assumption `validate_and_commit` increments version.

module-system/sov-state/src/storage.rs

@@ -169,6 +169,9 @@ pub trait Storage: Clone {
         + BorshSerialize
         + BorshDeserialize;
 
+    /// State update that will be committed to the database.
+    type StateUpdate;
+
     fn with_config(config: Self::RuntimeConfig) -> Result<Self, anyhow::Error>;
 
     /// Returns the value corresponding to the key or None if key is absent.
@@ -177,13 +180,30 @@ pub trait Storage: Clone {
     /// Returns the latest state root hash from the storage.
     fn get_state_root(&self, witness: &Self::Witness) -> anyhow::Result<[u8; 32]>;
 
+    /// Calculates new state root but does not commit any changes to the database.
+    fn compute_state_update(
+        &self,
+        state_accesses: OrderedReadsAndWrites,
+        witness: &Self::Witness,
+    ) -> Result<([u8; 32], Self::StateUpdate), anyhow::Error>;
+
+    /// Commits state changes to the database.
+    fn commit(&self, node_batch: &Self::StateUpdate);
+
     /// Validate all of the storage accesses in a particular cache log,
-    /// returning the new state root after applying all writes
+    /// returning the new state root after applying all writes.
+    /// This function is equivalent to calling:
+    /// `self.compute_state_update & self.commit`
     fn validate_and_commit(
         &self,
         state_accesses: OrderedReadsAndWrites,
         witness: &Self::Witness,
-    ) -> Result<[u8; 32], anyhow::Error>;
+    ) -> Result<[u8; 32], anyhow::Error> {
+        let (root_hash, node_batch) = self.compute_state_update(state_accesses, witness)?;
+        self.commit(&node_batch);
+
+        Ok(root_hash)
+    }
 
     /// Opens a storage access proof and validates it against a state root.
     /// It returns a result with the opened leaf (key, value) pair in case of success.

module-system/sov-state/src/zk_storage.rs

@@ -1,6 +1,7 @@
 use std::marker::PhantomData;
 use std::sync::Arc;
 
+use jmt::storage::NodeBatch;
 use jmt::{JellyfishMerkleTree, KeyHash, Version};
 #[cfg(all(target_os = "zkvm", feature = "bench"))]
 use zk_cycle_macros::cycle_tracker;
@@ -38,10 +39,9 @@ impl<S: MerkleProofSpec> ZkStorage<S> {
 
 impl<S: MerkleProofSpec> Storage for ZkStorage<S> {
     type Witness = S::Witness;
-
     type RuntimeConfig = [u8; 32];
-
     type Proof = jmt::proof::SparseMerkleProof<S::Hasher>;
+    type StateUpdate = NodeBatch;
 
     fn with_config(config: Self::RuntimeConfig) -> Result<Self, anyhow::Error> {
         Ok(Self::new(config))
@@ -56,11 +56,11 @@ impl<S: MerkleProofSpec> Storage for ZkStorage<S> {
     }
 
     #[cfg_attr(all(target_os = "zkvm", feature = "bench"), cycle_tracker)]
-    fn validate_and_commit(
+    fn compute_state_update(
         &self,
         state_accesses: OrderedReadsAndWrites,
         witness: &Self::Witness,
-    ) -> Result<[u8; 32], anyhow::Error> {
+    ) -> Result<([u8; 32], Self::StateUpdate), anyhow::Error> {
         let latest_version: Version = witness.get_hint();
         let reader = TreeWitnessReader::new(witness);
 
@@ -96,12 +96,15 @@ impl<S: MerkleProofSpec> Storage for ZkStorage<S> {
         // because the TreeReader is trusted
         let jmt = JellyfishMerkleTree::<_, S::Hasher>::new(&reader);
 
-        let (new_root, _tree_update) = jmt
+        let (new_root, tree_update) = jmt
             .put_value_set(batch, next_version)
             .expect("JMT update must succeed");
-        Ok(new_root.0)
+        Ok((new_root.0, tree_update.node_batch))
     }
 
+    #[cfg_attr(all(target_os = "zkvm", feature = "bench"), cycle_tracker)]
+    fn commit(&self, _node_batch: &Self::StateUpdate) {}
+
     fn is_empty(&self) -> bool {
         unimplemented!("Needs simplification in JellyfishMerkleTree: https://github.com/Sovereign-Labs/sovereign-sdk/issues/362")
     }