[WIP] Actually return & use StDiff

crates/batch-prover/src/da_block_handler.rs

@@ -20,6 +20,7 @@ use sov_prover_storage_manager::ProverStorageManager;
 use sov_rollup_interface::da::{BlockHeaderTrait, SequencerCommitment};
 use sov_rollup_interface::services::da::{DaService, SlotData};
 use sov_rollup_interface::spec::SpecId;
+use sov_rollup_interface::stateful_statediff::StatefulStateDiff;
 use sov_rollup_interface::zk::ZkvmHost;
 use tokio::select;
 use tokio::sync::Mutex;
@@ -306,8 +307,10 @@ pub(crate) fn break_sequencer_commitments_into_groups<DB: BatchProverLedgerOps>(
 
     let mut range = 0usize..=0usize;
     let mut cumulative_state_diff = StateDiff::new();
+    let mut cumulative_st_statediff = StatefulStateDiff::default();
     for (index, sequencer_commitment) in sequencer_commitments.iter().enumerate() {
         let mut sequencer_commitment_state_diff = StateDiff::new();
+        let mut sequencer_st_statediff = StatefulStateDiff::default();
         for l2_height in
             sequencer_commitment.l2_start_block_number..=sequencer_commitment.l2_end_block_number
         {
@@ -319,14 +322,27 @@ pub(crate) fn break_sequencer_commitments_into_groups<DB: BatchProverLedgerOps>(
                     sequencer_commitment.l2_end_block_number
                 ))?;
             sequencer_commitment_state_diff =
-                merge_state_diffs(sequencer_commitment_state_diff, state_diff);
+                merge_state_diffs(sequencer_commitment_state_diff, state_diff.0);
+            sequencer_st_statediff = sequencer_st_statediff.merge(state_diff.1);
         }
         cumulative_state_diff = merge_state_diffs(
             cumulative_state_diff,
             sequencer_commitment_state_diff.clone(),
         );
-
-        let compressed_state_diff = compress_blob(&borsh::to_vec(&cumulative_state_diff)?);
+        cumulative_st_statediff = cumulative_st_statediff.merge(sequencer_st_statediff.clone());
+
+        let uncompressed_state_diff = borsh::to_vec(&cumulative_state_diff)?;
+        let compressed_state_diff = compress_blob(&uncompressed_state_diff);
+        dbg!(uncompressed_state_diff.len(), compressed_state_diff.len());
+
+        let uncompressed_st_state_diff = borsh::to_vec(&cumulative_st_statediff)?;
+        let hex_state_diff = hex::encode(&uncompressed_st_state_diff);
+        let compressed_st_state_diff = compress_blob(&uncompressed_st_state_diff);
+        dbg!(
+            uncompressed_st_state_diff.len(),
+            compressed_st_state_diff.len(),
+            hex_state_diff
+        );
 
         // Threshold is checked by comparing compressed state diff size as the data will be compressed before it is written on DA
         let state_diff_threshold_reached = compressed_state_diff.len() > MAX_TXBODY_SIZE;
@@ -344,6 +360,7 @@ pub(crate) fn break_sequencer_commitments_into_groups<DB: BatchProverLedgerOps>(
             result_range.push(range);
             // Reset the cumulative state diff to be equal to the current commitment state diff
             cumulative_state_diff = sequencer_commitment_state_diff;
+            cumulative_st_statediff = sequencer_st_statediff;
             range = index..=index;
             current_spec = commitment_spec
         } else {

crates/sequencer/src/runner.rs

@@ -48,6 +48,7 @@ use sov_rollup_interface::da::{BlockHeaderTrait, DaSpec};
 use sov_rollup_interface::fork::ForkManager;
 use sov_rollup_interface::services::da::DaService;
 use sov_rollup_interface::soft_confirmation::{L2Header, SignedL2Header};
+use sov_rollup_interface::stateful_statediff::StatefulStateDiff;
 use sov_rollup_interface::zk::StorageRootHash;
 use sov_state::storage::NativeStorage;
 use sov_state::ProverStorage;
@@ -384,7 +385,7 @@ where
         da_block: <Da as DaService>::FilteredBlock,
         l1_fee_rate: u128,
         l2_block_mode: L2BlockMode,
-    ) -> anyhow::Result<(u64, u64, StateDiff)> {
+    ) -> anyhow::Result<(u64, u64, (StateDiff, StatefulStateDiff))> {
         let start = Instant::now();
         let da_height = da_block.header().height();
         let (l2_height, l1_height) = match self
@@ -769,7 +770,7 @@ where
                             // Only errors when there are no receivers
                             let _ = self.soft_confirmation_tx.send(l2_height);
 
-                            let _ = da_commitment_tx.send((l2_height, state_diff));
+                            let _ = da_commitment_tx.send((l2_height, state_diff.0));
                         },
                         Err(e) => {
                             error!("Sequencer error: {}", e);
@@ -801,7 +802,7 @@ where
                             // Only errors when there are no receivers
                             let _ = self.soft_confirmation_tx.send(l2_height);
 
-                            let _ = da_commitment_tx.send((l2_height, state_diff));
+                            let _ = da_commitment_tx.send((l2_height, state_diff.0));
                         },
                         Err(e) => {
                             error!("Sequencer error: {}", e);

crates/sovereign-sdk/full-node/db/sov-db/src/ledger_db/mod.rs

@@ -6,6 +6,7 @@ use rocksdb::WriteBatch;
 use sov_rollup_interface::da::SequencerCommitment;
 use sov_rollup_interface::fork::{Fork, ForkMigration};
 use sov_rollup_interface::soft_confirmation::L2Block;
+use sov_rollup_interface::stateful_statediff::StatefulStateDiff;
 use sov_rollup_interface::stf::StateDiff;
 use sov_rollup_interface::zk::{Proof, StorageRootHash};
 use sov_schema_db::{Schema, SchemaBatch, SeekKeyEncoder, DB};
@@ -574,7 +575,7 @@ impl BatchProverLedgerOps for LedgerDB {
     fn set_l2_state_diff(
         &self,
         l2_height: SoftConfirmationNumber,
-        state_diff: StateDiff,
+        state_diff: (StateDiff, StatefulStateDiff),
     ) -> anyhow::Result<()> {
         let mut schema_batch = SchemaBatch::new();
         schema_batch.put::<ProverStateDiffs>(&l2_height, &state_diff)?;
@@ -587,7 +588,7 @@ impl BatchProverLedgerOps for LedgerDB {
     fn get_l2_state_diff(
         &self,
         l2_height: SoftConfirmationNumber,
-    ) -> anyhow::Result<Option<StateDiff>> {
+    ) -> anyhow::Result<Option<(StateDiff, StatefulStateDiff)>> {
         self.db.get::<ProverStateDiffs>(&l2_height)
     }
 

crates/sovereign-sdk/full-node/db/sov-db/src/ledger_db/traits.rs

@@ -5,6 +5,7 @@ use anyhow::Result;
 use borsh::BorshSerialize;
 use sov_rollup_interface::da::SequencerCommitment;
 use sov_rollup_interface::soft_confirmation::L2Block;
+use sov_rollup_interface::stateful_statediff::StatefulStateDiff;
 use sov_rollup_interface::stf::StateDiff;
 use sov_rollup_interface::zk::{Proof, StorageRootHash};
 use sov_schema_db::SchemaBatch;
@@ -186,11 +187,14 @@ pub trait BatchProverLedgerOps: SharedLedgerOps + Send + Sync {
     fn set_l2_state_diff(
         &self,
         l2_height: SoftConfirmationNumber,
-        state_diff: StateDiff,
+        state_diff: (StateDiff, StatefulStateDiff),
     ) -> Result<()>;
 
     /// Returns an L2 state diff
-    fn get_l2_state_diff(&self, l2_height: SoftConfirmationNumber) -> Result<Option<StateDiff>>;
+    fn get_l2_state_diff(
+        &self,
+        l2_height: SoftConfirmationNumber,
+    ) -> Result<Option<(StateDiff, StatefulStateDiff)>>;
 
     /// Clears all pending proving sessions
     fn clear_pending_proving_sessions(&self) -> Result<()>;

crates/sovereign-sdk/full-node/db/sov-db/src/schema/tables.rs

@@ -14,6 +14,7 @@ use jmt::storage::{NibblePath, Node, NodeKey, StaleNodeIndex};
 use jmt::Version;
 use sov_rollup_interface::da::SequencerCommitment;
 use sov_rollup_interface::mmr::{MMRChunk, MMRNodeHash, Wtxid};
+use sov_rollup_interface::stateful_statediff::StatefulStateDiff;
 use sov_rollup_interface::stf::StateDiff;
 use sov_schema_db::schema::{KeyDecoder, KeyEncoder, ValueCodec};
 use sov_schema_db::{CodecError, SeekKeyEncoder};
@@ -436,7 +437,7 @@ define_table_with_default_codec!(
 
 define_table_with_default_codec!(
     /// L2 height to state diff for prover
-    (ProverStateDiffs) SoftConfirmationNumber => StateDiff
+    (ProverStateDiffs) SoftConfirmationNumber => (StateDiff, StatefulStateDiff)
 );
 
 define_table_with_seek_key_codec!(

crates/sovereign-sdk/full-node/db/sov-db/src/schema/types/batch_proof.rs

@@ -124,7 +124,7 @@ impl From<StoredBatchProofOutput> for BatchProofOutputRpcResponse {
             StoredBatchProofOutput::V3(value) => Self {
                 initial_state_root: value.initial_state_root.to_vec(),
                 final_state_root: value.final_state_root.to_vec(),
-                state_diff: value.state_diff,
+                state_diff: value.state_diff.0,
                 da_slot_hash: None,
                 sequencer_da_public_key: vec![],
                 sequencer_public_key: vec![],

crates/sovereign-sdk/module-system/sov-modules-core/Cargo.toml

@@ -13,7 +13,9 @@ resolver = "2"
 
 
 [dependencies]
+alloy-primitives = { workspace = true, default-features = false, features = ["serde"] }
 anyhow = { workspace = true }
+bcs = { workspace = true }
 bech32 = { workspace = true }
 borsh = { workspace = true }
 derive_more = { workspace = true, features = ["display", "into"] }

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

@@ -7,6 +7,7 @@ use core::fmt;
 use borsh::{BorshDeserialize, BorshSerialize};
 #[cfg(feature = "sync")]
 use serde::Serialize;
+use sov_rollup_interface::stateful_statediff::StatefulStateDiff;
 use sov_rollup_interface::stf::{StateDiff, StateRootTransition};
 use sov_rollup_interface::witness::Witness;
 use sov_rollup_interface::zk::{SparseMerkleProofSha2, StorageRootHash};
@@ -209,7 +210,7 @@ pub trait Storage: Clone {
         (
             StateRootTransition,
             Self::StateUpdate,
-            StateDiff, // computed in Zk mode
+            (StateDiff, StatefulStateDiff), // computed in Zk mode
         ),
         anyhow::Error,
     >;

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

@@ -24,6 +24,7 @@ use sov_rollup_interface::da::SequencerCommitment;
 use sov_rollup_interface::fork::ForkManager;
 use sov_rollup_interface::soft_confirmation::{L2Block, SignedL2Header};
 use sov_rollup_interface::spec::SpecId;
+use sov_rollup_interface::stateful_statediff::StatefulStateDiff;
 use sov_rollup_interface::stf::{
     SoftConfirmationError, SoftConfirmationResult, StateTransitionError,
 };
@@ -87,7 +88,7 @@ pub struct ApplySequencerCommitmentsOutput {
     /// Final state root after all sequencer commitments were applied
     pub final_state_root: StorageRootHash,
     /// State diff generated after applying
-    pub state_diff: CumulativeStateDiff,
+    pub state_diff: (CumulativeStateDiff, StatefulStateDiff),
     /// Last processed L2 block height
     pub last_l2_height: u64,
     /// Last soft confirmation hash
@@ -424,6 +425,7 @@ where
         forks: &[Fork],
     ) -> ApplySequencerCommitmentsOutput {
         let mut state_diff = CumulativeStateDiff::default();
+        let mut st_statediff = StatefulStateDiff::default();
 
         let sequencer_commitment_merkle_roots = sequencer_commitments
             .iter()
@@ -570,7 +572,8 @@ where
 
                 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);
+                state_diff.extend(result.state_diff.0);
+                st_statediff = st_statediff.merge(result.state_diff.1);
 
                 l2_height += 1;
 
@@ -619,7 +622,7 @@ where
 
         ApplySequencerCommitmentsOutput {
             final_state_root: current_state_root,
-            state_diff,
+            state_diff: (state_diff, st_statediff),
             // There has to be a height
             last_l2_height: last_commitment_end_height.unwrap(),
             final_soft_confirmation_hash: prev_soft_confirmation_hash.unwrap(),

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

@@ -7,9 +7,6 @@ pub mod codec;
 #[cfg(feature = "native")]
 mod prover_storage;
 
-/// Stateful Statediff primitives
-pub mod stateful_statediff;
-
 mod zk_storage;
 
 #[cfg(feature = "native")]

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

@@ -10,6 +10,7 @@ use sov_modules_core::{
     CacheKey, NativeStorage, OrderedWrites, ReadWriteLog, Storage, StorageKey, StorageProof,
     StorageValue,
 };
+use sov_rollup_interface::stateful_statediff::{self, StatefulStateDiff};
 use sov_rollup_interface::stf::{StateDiff, StateRootTransition};
 use sov_rollup_interface::witness::Witness;
 use sov_rollup_interface::zk::StorageRootHash;
@@ -139,7 +140,14 @@ impl Storage for ProverStorage {
         &self,
         state_log: &ReadWriteLog,
         witness: &mut Witness,
-    ) -> Result<(StateRootTransition, Self::StateUpdate, StateDiff), anyhow::Error> {
+    ) -> Result<
+        (
+            StateRootTransition,
+            Self::StateUpdate,
+            (StateDiff, StatefulStateDiff),
+        ),
+        anyhow::Error,
+    > {
         let version = self.version();
         let jmt = JellyfishMerkleTree::<_, DefaultHasher>::new(&self.db);
 
@@ -170,11 +178,19 @@ impl Storage for ProverStorage {
             witness.add_hint(&proof);
         }
 
-        let pre_state = crate::stateful_statediff::build_pre_state(state_log.ordered_reads());
+        let pre_state =
+            stateful_statediff::build_pre_state(state_log.ordered_reads().iter().map(|(k, v)| {
+                let k = k.key.clone();
+                let v = v.as_ref().map(|v| v.value.clone());
+                (k, v)
+            }));
         let post_state =
-            crate::stateful_statediff::build_post_state(state_log.iter_ordered_writes());
-
-        let _st_statediff = crate::stateful_statediff::compress_state(pre_state, post_state);
+            stateful_statediff::build_post_state(state_log.iter_ordered_writes().map(|(k, v)| {
+                let k = k.key.clone();
+                let v = v.as_ref().map(|v| v.value.clone());
+                (k, v)
+            }));
+        let st_statediff = stateful_statediff::compress_state(pre_state, post_state);
 
         let mut key_preimages = vec![];
         let mut diff = vec![];
@@ -198,15 +214,15 @@ impl Storage for ProverStorage {
             .put_value_set_with_proof(batch, next_version)
             .expect("JMT update must succeed");
 
-        let unparsed_len: usize = _st_statediff
+        let unparsed_len: usize = st_statediff
             .unparsed
             .iter()
             .map(|(_k, v)| if let Some(x) = v { x.len() } else { 0 })
             .sum();
-        let ststdiff = borsh::to_vec(&_st_statediff).unwrap();
-        let _orig: crate::stateful_statediff::StatefulStateDiff =
-            borsh::from_slice(&ststdiff).unwrap(); // check if we can parse it
+        let ststdiff = borsh::to_vec(&st_statediff).unwrap();
+        let _orig: StatefulStateDiff = borsh::from_slice(&ststdiff).unwrap(); // check if we can parse it
         let prevdiff = borsh::to_vec(&diff).unwrap();
+        let _ = st_statediff;
 
         println!(
             "ststdiff: {} bytes, diff: {} bytes, ststdiff unparsed: {} bytes \n",
@@ -233,7 +249,7 @@ impl Storage for ProverStorage {
                 final_root: new_root.into(),
             },
             state_update,
-            diff,
+            (diff, st_statediff),
         ))
     }