Create a wrapper struct for state root transition

crates/batch-prover/src/runner.rs

@@ -401,7 +401,7 @@ where
 
         let receipt = soft_confirmation_result.soft_confirmation_receipt;
 
-        let next_state_root = soft_confirmation_result.final_state_root;
+        let next_state_root = soft_confirmation_result.state_root_transition.final_root;
         // Check if post state root is the same as the one in the soft confirmation
         if next_state_root.as_ref().to_vec() != soft_confirmation.state_root {
             bail!("Post state root mismatch at height: {}", l2_height)

crates/evm/src/tests/utils.rs

@@ -100,7 +100,7 @@ pub(crate) fn commit(
 
     let (cache_log, mut witness) = checkpoint.freeze();
 
-    let ((_, root), authenticated_node_batch, _) = storage
+    let (state_root_transition, authenticated_node_batch, _) = storage
         .compute_state_update(cache_log, &mut witness)
         .expect("jellyfish merkle tree update must succeed");
 
@@ -111,7 +111,7 @@ pub(crate) fn commit(
 
     storage.commit(&authenticated_node_batch, &accessory_log, &offchain_log);
 
-    root.0
+    state_root_transition.final_root.0
 }
 
 /// Loads the genesis configuration from the given path and pushes the accounts to the evm config

crates/fullnode/src/runner.rs

@@ -260,7 +260,7 @@ where
 
         let receipt = soft_confirmation_result.soft_confirmation_receipt;
 
-        let next_state_root = soft_confirmation_result.final_state_root;
+        let next_state_root = soft_confirmation_result.state_root_transition.final_root;
         // Check if post state root is the same as the one in the soft confirmation
         if next_state_root.as_ref().to_vec() != soft_confirmation.state_root {
             bail!("Post state root mismatch at height: {}", l2_height)

crates/fullnode/tests/hash_stf.rs

@@ -43,7 +43,7 @@ impl HashStf {
             ordered_writes: vec![(hash_key.to_cache_key(), Some(hash_value.into_cache_value()))],
         };
 
-        let ((_, jmt_root_hash), state_update, _) = storage
+        let (state_root_transition, state_update, _) = storage
             .compute_state_update(ordered_reads_writes, witness)
             .unwrap();
 
@@ -55,7 +55,13 @@ impl HashStf {
 
         let mut root_hash = [0u8; 32];
 
-        for (i, &byte) in jmt_root_hash.as_ref().iter().enumerate().take(32) {
+        for (i, &byte) in state_root_transition
+            .final_root
+            .as_ref()
+            .iter()
+            .enumerate()
+            .take(32)
+        {
             root_hash[i] = byte;
         }
 

crates/sequencer/src/runner.rs

@@ -492,19 +492,20 @@ where
                     prestate,
                     &mut signed_soft_confirmation,
                 );
+                let state_root_transition = soft_confirmation_result.state_root_transition;
 
                 let receipt = soft_confirmation_result.soft_confirmation_receipt;
 
-                if soft_confirmation_result.final_state_root.as_ref() == self.state_root.as_ref() {
+                if state_root_transition.final_root.as_ref() == self.state_root.as_ref() {
                     bail!("Max L2 blocks per L1 is reached for the current L1 block. State root is the same as before, skipping");
                 }
 
                 trace!(
                     "State root after applying slot: {:?}",
-                    soft_confirmation_result.final_state_root
+                    state_root_transition.final_root,
                 );
 
-                let next_state_root = soft_confirmation_result.final_state_root;
+                let next_state_root = state_root_transition.final_root;
 
                 self.storage_manager
                     .save_change_set_l2(l2_height, soft_confirmation_result.change_set)?;

crates/sovereign-sdk/module-system/sov-modules-core/src/storage/mod.rs

@@ -6,7 +6,7 @@ use core::fmt;
 use borsh::{BorshDeserialize, BorshSerialize};
 use serde::de::DeserializeOwned;
 use serde::Serialize;
-use sov_rollup_interface::stf::StateDiff;
+use sov_rollup_interface::stf::{StateDiff, StateRootTransition};
 use sov_rollup_interface::RefCount;
 
 use crate::common::{AlignedVec, Prefix, Version, Witness};
@@ -239,13 +239,14 @@ pub trait Storage: Clone {
     }
 
     /// Calculates new state root but does not commit any changes to the database.
+    #[allow(clippy::type_complexity)]
     fn compute_state_update(
         &self,
         state_accesses: OrderedReadsAndWrites,
         witness: &mut Self::Witness,
     ) -> Result<
         (
-            (Self::Root, Self::Root),
+            StateRootTransition<Self::Root>,
             Self::StateUpdate,
             StateDiff, // computed in Zk mode
         ),
@@ -268,10 +269,11 @@ pub trait Storage: Clone {
         accessory_update: &OrderedReadsAndWrites,
         offchain_update: &OrderedReadsAndWrites,
     ) -> Result<Self::Root, anyhow::Error> {
-        let ((_, root_hash), node_batch, _) = self.compute_state_update(state_accesses, witness)?;
+        let (state_root_transition, node_batch, _) =
+            self.compute_state_update(state_accesses, witness)?;
         self.commit(&node_batch, accessory_update, offchain_update);
 
-        Ok(root_hash)
+        Ok(state_root_transition.final_root)
     }
 
     /// Validate all of the storage accesses in a particular cache log,

crates/sovereign-sdk/module-system/sov-modules-stf-blueprint/src/lib.rs

@@ -345,21 +345,23 @@ where
             );
         }
 
-        let (init_state_root, final_state_root, witness, offchain_witness, storage, state_diff) = {
+        let (state_root_transition, witness, offchain_witness, storage, state_diff) = {
             let working_set = checkpoint.to_revertable();
             // Save checkpoint
             let mut checkpoint = working_set.checkpoint();
 
             let (cache_log, mut witness) = checkpoint.freeze();
 
-            let ((init_root_hash, final_root_hash), state_update, state_diff) = pre_state
+            let (state_root_transition, state_update, state_diff) = pre_state
                 .compute_state_update(cache_log, &mut witness)
                 .expect("jellyfish merkle tree update must succeed");
 
             let mut working_set = checkpoint.to_revertable();
 
-            self.runtime
-                .finalize_hook(&final_root_hash, &mut working_set.accessory_state());
+            self.runtime.finalize_hook(
+                &state_root_transition.final_root,
+                &mut working_set.accessory_state(),
+            );
 
             let mut checkpoint = working_set.checkpoint();
             let accessory_log = checkpoint.freeze_non_provable();
@@ -368,8 +370,7 @@ where
             pre_state.commit(&state_update, &accessory_log, &offchain_log);
 
             (
-                init_root_hash,
-                final_root_hash,
+                state_root_transition,
                 witness,
                 offchain_witness,
                 pre_state,
@@ -378,8 +379,7 @@ where
         };
 
         SoftConfirmationResult {
-            init_state_root,
-            final_state_root,
+            state_root_transition,
             change_set: storage,
             witness,
             offchain_witness,
@@ -422,9 +422,10 @@ where
         let mut checkpoint = working_set.checkpoint();
         let (log, mut witness) = checkpoint.freeze();
 
-        let ((_, genesis_hash), state_update, _) = pre_state
+        let (init_and_final_roots, state_update, _) = pre_state
             .compute_state_update(log, &mut witness)
             .expect("Storage update must succeed");
+        let genesis_hash = init_and_final_roots.final_root;
 
         let mut working_set = checkpoint.to_revertable();
 
@@ -809,8 +810,8 @@ where
                     // for now we don't allow "broken" seq. com.s
                     .expect("Soft confirmation must succeed");
 
-                assert_eq!(current_state_root, result.init_state_root);
-                current_state_root = result.final_state_root;
+                assert_eq!(current_state_root, result.state_root_transition.init_root);
+                current_state_root = result.state_root_transition.final_root;
                 state_diff.extend(result.state_diff);
 
                 // Notify fork manager about the block so that the next spec / fork

crates/sovereign-sdk/module-system/sov-state/src/prover_storage.rs

@@ -9,7 +9,7 @@ use sov_modules_core::{
     CacheKey, NativeStorage, OrderedReadsAndWrites, Storage, StorageKey, StorageProof,
     StorageValue, Witness,
 };
-use sov_rollup_interface::stf::StateDiff;
+use sov_rollup_interface::stf::{StateDiff, StateRootTransition};
 
 use crate::config::Config;
 use crate::{DefaultHasher, DefaultWitness};
@@ -119,7 +119,14 @@ where
         &self,
         state_accesses: OrderedReadsAndWrites,
         witness: &mut Self::Witness,
-    ) -> Result<((Self::Root, Self::Root), Self::StateUpdate, StateDiff), anyhow::Error> {
+    ) -> Result<
+        (
+            StateRootTransition<Self::Root>,
+            Self::StateUpdate,
+            StateDiff,
+        ),
+        anyhow::Error,
+    > {
         let latest_version = self.db.get_next_version() - 1;
         let jmt = JellyfishMerkleTree::<_, DefaultHasher>::new(&self.db);
 
@@ -191,7 +198,14 @@ where
         // We need the state diff to be calculated only inside zk context.
         // The diff then can be used by special nodes to construct the state of the rollup by verifying the zk proof.
         // And constructing the tree from the diff.
-        Ok(((prev_root, new_root), state_update, diff))
+        Ok((
+            StateRootTransition {
+                init_root: prev_root,
+                final_root: new_root,
+            },
+            state_update,
+            diff,
+        ))
     }
 
     fn commit(

crates/sovereign-sdk/module-system/sov-state/src/zk_storage.rs

@@ -6,7 +6,7 @@ use sha2::Digest;
 use sov_modules_core::{
     OrderedReadsAndWrites, Storage, StorageKey, StorageProof, StorageValue, Witness,
 };
-use sov_rollup_interface::stf::StateDiff;
+use sov_rollup_interface::stf::{StateDiff, StateRootTransition};
 
 /// A [`Storage`] implementation designed to be used inside the zkVM.
 #[derive(Default)]
@@ -76,7 +76,14 @@ where
         &self,
         state_accesses: OrderedReadsAndWrites,
         witness: &mut Self::Witness,
-    ) -> Result<((Self::Root, Self::Root), Self::StateUpdate, StateDiff), anyhow::Error> {
+    ) -> Result<
+        (
+            StateRootTransition<Self::Root>,
+            Self::StateUpdate,
+            StateDiff,
+        ),
+        anyhow::Error,
+    > {
         let prev_state_root = witness.get_hint();
 
         // For each value that's been read from the tree, verify the provided smt proof
@@ -124,7 +131,10 @@ where
             .expect("Updates must be valid");
 
         Ok((
-            (jmt::RootHash(prev_state_root), jmt::RootHash(new_root)),
+            StateRootTransition {
+                init_root: jmt::RootHash(prev_state_root),
+                final_root: jmt::RootHash(new_root),
+            },
             (),
             diff,
         ))

crates/sovereign-sdk/rollup-interface/src/state_machine/stf.rs

@@ -143,6 +143,14 @@ pub struct SlotResult<S, Cs, B, T, W> {
     pub state_diff: StateDiff,
 }
 
+/// Helper struct which contains initial and final state roots.
+pub struct StateRootTransition<Root> {
+    /// Initial state root
+    pub init_root: Root,
+    /// Final state root
+    pub final_root: Root,
+}
+
 /// Result of applying a soft confirmation to current state
 /// Where:
 /// - S - generic for state root
@@ -151,10 +159,8 @@ pub struct SlotResult<S, Cs, B, T, W> {
 /// - W - generic for witness
 /// - Da - generic for DA layer
 pub struct SoftConfirmationResult<S, Cs, T, W, Da: DaSpec> {
-    /// Initial state root read from the hints
-    pub init_state_root: S,
-    /// Finals state root after all soft confirmation txs are applied
-    pub final_state_root: S,
+    /// Contains state root before and after applying txs
+    pub state_root_transition: StateRootTransition<S>,
     /// Container for all state alterations that happened during soft confirmation execution
     pub change_set: Cs,
     /// Witness after applying the whole block