Move StDiff into rollup-interface

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-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;
@@ -152,11 +153,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![];
@@ -180,15 +189,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",

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

@@ -2,6 +2,7 @@ use jmt::KeyHash;
 use sov_modules_core::{
     OrderedWrites, ReadWriteLog, Storage, StorageKey, StorageProof, StorageValue,
 };
+use sov_rollup_interface::stateful_statediff;
 use sov_rollup_interface::stf::{StateDiff, StateRootTransition};
 use sov_rollup_interface::witness::Witness;
 use sov_rollup_interface::zk::StorageRootHash;
@@ -47,11 +48,19 @@ impl Storage for ZkStorage {
             proof.verify(jmt::RootHash(prev_state_root), key_hash, value)?;
         }
 
-        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 diff = vec![];
 
@@ -81,13 +90,14 @@ impl Storage for ZkStorage {
             )
             .expect("Updates must be valid");
 
-        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 ststdiff = borsh::to_vec(&st_statediff).unwrap();
         let prevdiff = borsh::to_vec(&diff).unwrap();
+        let _ = st_statediff;
 
         println!(
             "zk: ststdiff: {} bytes, diff: {} bytes, ststdiff unparsed: {} bytes \n",

crates/sovereign-sdk/rollup-interface/Cargo.toml

@@ -17,7 +17,8 @@ resolver = "2"
 [dependencies]
 anyhow = { workspace = true, features = ["default"]}
 async-trait = { workspace = true, optional = true }
-alloy-primitives = { workspace = true, features=["serde"], optional = true}
+alloy-primitives = { workspace = true, features=["serde"]}
+bcs = { workspace = true }
 borsh = { workspace = true, features = ["default", "bytes"]}
 bytes = { workspace = true,  default-features = true }
 digest = { workspace = true, features = ["default"]}
@@ -38,5 +39,5 @@ serde_json = { workspace = true }
 
 [features]
 default = []
-native = ["alloy-primitives", "tokio", "futures", "tracing", "risc0-zkp", "faster-hex", "async-trait", "thiserror", "hex"]
+native = ["tokio", "futures", "tracing", "risc0-zkp", "faster-hex", "async-trait", "thiserror", "hex"]
 testing = []

crates/sovereign-sdk/rollup-interface/src/lib.rs

@@ -19,6 +19,8 @@ mod node;
 /// Specs module
 pub mod spec;
 mod state_machine;
+/// StatefulStateDiff module
+pub mod stateful_statediff;
 
 #[cfg(not(feature = "native"))]
 pub use std::rc::Rc as RefCount;

crates/sovereign-sdk/module-system/sov-state/src/stateful_statediff/mod.rs → crates/sovereign-sdk/rollup-interface/src/stateful_statediff/mod.rs

@@ -1,3 +1,5 @@
+//! Stateful StateDiff construction and compression
+
 /// Compression primitives
 pub mod compression;
 
@@ -8,10 +10,11 @@ use alloy_primitives::{Address, B256, U256};
 use borsh::{BorshDeserialize, BorshSerialize};
 use compression::{CodeHashChange, SlotChange};
 use serde::{Deserialize, Serialize};
-use sov_modules_core::{CacheKey, CacheValue};
-use sov_rollup_interface::RefCount;
 
-pub(crate) struct PreState {
+use crate::RefCount;
+
+/// Reflects state before applying changes
+pub struct PreState {
     evm_accounts_prefork2: BTreeMap<Address, Option<DbAccountInfo>>,
     evm_storage_prefork2: BTreeMap<Address, BTreeMap<U256, Option<U256>>>,
     evm_accounts: BTreeMap<u64, Option<DbAccountInfo>>,
@@ -23,7 +26,10 @@ fn borsh_u256_from_slice(v: impl AsRef<[u8]>) -> U256 {
     U256::from_limbs(s)
 }
 
-pub(crate) fn build_pre_state(ordered_reads: &[(CacheKey, Option<CacheValue>)]) -> PreState {
+/// Create a PreState which reflects the state before applying changes
+pub fn build_pre_state(
+    ordered_reads: impl Iterator<Item = (RefCount<[u8]>, Option<RefCount<[u8]>>)>,
+) -> PreState {
     // We need the first values we read. So we traverse from the beginning.
     // We are only interested in keys -> values only when we see them the first time.
     // And we need only Evm accounts and storage, because that's the only
@@ -33,8 +39,8 @@ pub(crate) fn build_pre_state(ordered_reads: &[(CacheKey, Option<CacheValue>)])
     let mut evm_accounts = BTreeMap::new();
     let mut evm_storage = BTreeMap::new();
 
-    for (k, v) in ordered_reads {
-        let (key, value) = (k.key.as_ref(), v.as_ref().map(|v| v.value.as_ref()));
+    for (cache_key, cache_value) in ordered_reads {
+        let (key, value) = (cache_key.as_ref(), cache_value.as_ref().map(|v| v.as_ref()));
         match &key[..6] {
             _account_prefork2 @ b"Evm/a/" => {
                 let address: Address = bcs::from_bytes(&key[6..]).unwrap();
@@ -84,7 +90,8 @@ pub(crate) fn build_pre_state(ordered_reads: &[(CacheKey, Option<CacheValue>)])
 /// A diff of the state, represented as a list of key-value pairs.
 pub type UnparsedStateDiff = Vec<(RefCount<[u8]>, Option<RefCount<[u8]>>)>;
 
-pub(crate) struct PostState {
+/// Reflects state after applying changes
+pub struct PostState {
     evm_accounts_prefork2: BTreeMap<Address, Option<DbAccountInfo>>,
     evm_storage_prefork2: BTreeMap<Address, BTreeMap<U256, Option<U256>>>,
     evm_accounts: BTreeMap<u64, Option<DbAccountInfo>>,
@@ -95,8 +102,9 @@ pub(crate) struct PostState {
     unparsed: UnparsedStateDiff,
 }
 
-pub(crate) fn build_post_state<'a>(
-    ordered_writes: impl Iterator<Item = (&'a CacheKey, &'a Option<CacheValue>)>,
+/// Create a PostState which reflects the state after applying changes
+pub fn build_post_state(
+    ordered_writes: impl Iterator<Item = (RefCount<[u8]>, Option<RefCount<[u8]>>)>,
 ) -> PostState {
     // We need the last values we write. So we traverse from the end.
     let mut evm_accounts_prefork2: BTreeMap<Address, Option<DbAccountInfo>> = BTreeMap::new();
@@ -109,10 +117,7 @@ pub(crate) fn build_post_state<'a>(
     let mut unparsed = UnparsedStateDiff::new();
 
     for (cache_key, cache_value) in ordered_writes.into_iter() {
-        let (key, value) = (
-            cache_key.key.as_ref(),
-            cache_value.as_ref().map(|v| v.value.as_ref()),
-        );
+        let (key, value) = (cache_key.as_ref(), cache_value.as_ref().map(|v| v.as_ref()));
         match &key[..6] {
             _account_prefork2 @ b"Evm/a/" => {
                 // Only the first key -> value
@@ -194,10 +199,8 @@ pub(crate) fn build_post_state<'a>(
 
                 // let hx_key = alloy_primitives::hex::encode(key);
                 // println!("unknown key: {}", hx_key);
-                let key_bytes = cache_key.key.clone();
-                let value_bytes = cache_value.as_ref().map(|v| v.value.clone());
 
-                unparsed.push((key_bytes, value_bytes));
+                unparsed.push((cache_key, cache_value));
             }
         }
     }
@@ -238,7 +241,7 @@ pub struct LatestBlockHashes {
 }
 
 /// Reflects all state change
-#[derive(Debug, BorshSerialize, BorshDeserialize)]
+#[derive(Debug, Default, BorshSerialize, BorshDeserialize)]
 pub struct StatefulStateDiff {
     // TODO: Remove before mainnet
     #[borsh(
@@ -274,7 +277,8 @@ pub struct StatefulStateDiff {
     pub unparsed: UnparsedStateDiff,
 }
 
-pub(crate) fn compress_state(pre_state: PreState, post_state: PostState) -> StatefulStateDiff {
+/// Create a StatefulStateDiff which reflects the state diff after applying changes
+pub fn compress_state(pre_state: PreState, post_state: PostState) -> StatefulStateDiff {
     use compression::{
         compress_one_best_strategy, compress_one_code_hash, compress_two_best_strategy,
         compress_two_code_hash,