From 1863460e6b4560c6e36b2b3e50ce0a57b48943d2 Mon Sep 17 00:00:00 2001 From: Roman Proskuryakoff Date: Sat, 22 Feb 2025 20:34:29 +0300 Subject: [PATCH] Move StDiff into rollup-interface --- Cargo.lock | 3 ++ .../module-system/sov-modules-core/Cargo.toml | 2 + .../module-system/sov-state/src/lib.rs | 3 -- .../sov-state/src/prover_storage.rs | 25 ++++++++---- .../module-system/sov-state/src/zk_storage.rs | 22 +++++++--- .../sovereign-sdk/rollup-interface/Cargo.toml | 5 ++- .../sovereign-sdk/rollup-interface/src/lib.rs | 2 + .../src/stateful_statediff/compression.rs | 0 .../src/stateful_statediff/mod.rs | 40 ++++++++++--------- guests/risc0/batch-proof/bitcoin/Cargo.lock | 4 ++ guests/risc0/batch-proof/mock/Cargo.lock | 4 ++ .../light-client-proof/bitcoin/Cargo.lock | 4 ++ .../risc0/light-client-proof/mock/Cargo.lock | 4 ++ 13 files changed, 81 insertions(+), 37 deletions(-) rename crates/sovereign-sdk/{module-system/sov-state => rollup-interface}/src/stateful_statediff/compression.rs (100%) rename crates/sovereign-sdk/{module-system/sov-state => rollup-interface}/src/stateful_statediff/mod.rs (95%) diff --git a/Cargo.lock b/Cargo.lock index c15d761a31..2be2e2ebc1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8001,7 +8001,9 @@ dependencies = [ name = "sov-modules-core" version = "0.6.0" dependencies = [ + "alloy-primitives", "anyhow", + "bcs", "bech32 0.9.1", "borsh", "derive_more", @@ -8102,6 +8104,7 @@ dependencies = [ "alloy-primitives", "anyhow", "async-trait", + "bcs", "borsh", "bytes", "digest 0.10.7", diff --git a/crates/sovereign-sdk/module-system/sov-modules-core/Cargo.toml b/crates/sovereign-sdk/module-system/sov-modules-core/Cargo.toml index aab42ddb4e..5b05375875 100644 --- a/crates/sovereign-sdk/module-system/sov-modules-core/Cargo.toml +++ b/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"] } diff --git a/crates/sovereign-sdk/module-system/sov-state/src/lib.rs b/crates/sovereign-sdk/module-system/sov-state/src/lib.rs index 38edfefb1b..c7eaeed080 100644 --- a/crates/sovereign-sdk/module-system/sov-state/src/lib.rs +++ b/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")] diff --git a/crates/sovereign-sdk/module-system/sov-state/src/prover_storage.rs b/crates/sovereign-sdk/module-system/sov-state/src/prover_storage.rs index c6f90e5190..93f2616306 100644 --- a/crates/sovereign-sdk/module-system/sov-state/src/prover_storage.rs +++ b/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", diff --git a/crates/sovereign-sdk/module-system/sov-state/src/zk_storage.rs b/crates/sovereign-sdk/module-system/sov-state/src/zk_storage.rs index 8188b2d52a..0f6f372f78 100644 --- a/crates/sovereign-sdk/module-system/sov-state/src/zk_storage.rs +++ b/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", diff --git a/crates/sovereign-sdk/rollup-interface/Cargo.toml b/crates/sovereign-sdk/rollup-interface/Cargo.toml index e0d8833e23..5be1a11e13 100644 --- a/crates/sovereign-sdk/rollup-interface/Cargo.toml +++ b/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 = [] diff --git a/crates/sovereign-sdk/rollup-interface/src/lib.rs b/crates/sovereign-sdk/rollup-interface/src/lib.rs index ba68876683..69fced847d 100644 --- a/crates/sovereign-sdk/rollup-interface/src/lib.rs +++ b/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; diff --git a/crates/sovereign-sdk/module-system/sov-state/src/stateful_statediff/compression.rs b/crates/sovereign-sdk/rollup-interface/src/stateful_statediff/compression.rs similarity index 100% rename from crates/sovereign-sdk/module-system/sov-state/src/stateful_statediff/compression.rs rename to crates/sovereign-sdk/rollup-interface/src/stateful_statediff/compression.rs diff --git a/crates/sovereign-sdk/module-system/sov-state/src/stateful_statediff/mod.rs b/crates/sovereign-sdk/rollup-interface/src/stateful_statediff/mod.rs similarity index 95% rename from crates/sovereign-sdk/module-system/sov-state/src/stateful_statediff/mod.rs rename to crates/sovereign-sdk/rollup-interface/src/stateful_statediff/mod.rs index e79a0c4f07..9c0dafb7a5 100644 --- a/crates/sovereign-sdk/module-system/sov-state/src/stateful_statediff/mod.rs +++ b/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>, evm_storage_prefork2: BTreeMap>>, evm_accounts: BTreeMap>, @@ -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)]) -> PreState { +/// Create a PreState which reflects the state before applying changes +pub fn build_pre_state( + ordered_reads: impl Iterator, Option>)>, +) -> 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)]) 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)]) /// A diff of the state, represented as a list of key-value pairs. pub type UnparsedStateDiff = Vec<(RefCount<[u8]>, Option>)>; -pub(crate) struct PostState { +/// Reflects state after applying changes +pub struct PostState { evm_accounts_prefork2: BTreeMap>, evm_storage_prefork2: BTreeMap>>, evm_accounts: BTreeMap>, @@ -95,8 +102,9 @@ pub(crate) struct PostState { unparsed: UnparsedStateDiff, } -pub(crate) fn build_post_state<'a>( - ordered_writes: impl Iterator)>, +/// Create a PostState which reflects the state after applying changes +pub fn build_post_state( + ordered_writes: impl Iterator, Option>)>, ) -> PostState { // We need the last values we write. So we traverse from the end. let mut evm_accounts_prefork2: BTreeMap> = 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, diff --git a/guests/risc0/batch-proof/bitcoin/Cargo.lock b/guests/risc0/batch-proof/bitcoin/Cargo.lock index c1705c2152..a5de660031 100644 --- a/guests/risc0/batch-proof/bitcoin/Cargo.lock +++ b/guests/risc0/batch-proof/bitcoin/Cargo.lock @@ -3268,7 +3268,9 @@ dependencies = [ name = "sov-modules-core" version = "0.6.0" dependencies = [ + "alloy-primitives", "anyhow", + "bcs", "bech32 0.9.1", "borsh", "derive_more", @@ -3315,7 +3317,9 @@ dependencies = [ name = "sov-rollup-interface" version = "0.6.0" dependencies = [ + "alloy-primitives", "anyhow", + "bcs", "borsh", "bytes", "digest 0.10.7", diff --git a/guests/risc0/batch-proof/mock/Cargo.lock b/guests/risc0/batch-proof/mock/Cargo.lock index 44ffa21b74..f56b89b80c 100644 --- a/guests/risc0/batch-proof/mock/Cargo.lock +++ b/guests/risc0/batch-proof/mock/Cargo.lock @@ -3171,7 +3171,9 @@ dependencies = [ name = "sov-modules-core" version = "0.6.0" dependencies = [ + "alloy-primitives", "anyhow", + "bcs", "bech32", "borsh", "derive_more", @@ -3218,7 +3220,9 @@ dependencies = [ name = "sov-rollup-interface" version = "0.6.0" dependencies = [ + "alloy-primitives", "anyhow", + "bcs", "borsh", "bytes", "digest 0.10.7", diff --git a/guests/risc0/light-client-proof/bitcoin/Cargo.lock b/guests/risc0/light-client-proof/bitcoin/Cargo.lock index 3ae1fc2aa0..3617c8fbee 100644 --- a/guests/risc0/light-client-proof/bitcoin/Cargo.lock +++ b/guests/risc0/light-client-proof/bitcoin/Cargo.lock @@ -2065,7 +2065,9 @@ dependencies = [ name = "sov-modules-core" version = "0.6.0" dependencies = [ + "alloy-primitives", "anyhow", + "bcs", "bech32 0.9.1", "borsh", "derive_more", @@ -2100,7 +2102,9 @@ dependencies = [ name = "sov-rollup-interface" version = "0.6.0" dependencies = [ + "alloy-primitives", "anyhow", + "bcs", "borsh", "bytes", "digest", diff --git a/guests/risc0/light-client-proof/mock/Cargo.lock b/guests/risc0/light-client-proof/mock/Cargo.lock index 95b0895ab0..16b1f0397d 100644 --- a/guests/risc0/light-client-proof/mock/Cargo.lock +++ b/guests/risc0/light-client-proof/mock/Cargo.lock @@ -1968,7 +1968,9 @@ dependencies = [ name = "sov-modules-core" version = "0.6.0" dependencies = [ + "alloy-primitives", "anyhow", + "bcs", "bech32", "borsh", "derive_more", @@ -2003,7 +2005,9 @@ dependencies = [ name = "sov-rollup-interface" version = "0.6.0" dependencies = [ + "alloy-primitives", "anyhow", + "bcs", "borsh", "bytes", "digest",