From cb8ed94ae2c72d3acd1e00887a50eee43a782038 Mon Sep 17 00:00:00 2001 From: marijamijailovic Date: Thu, 22 Aug 2024 12:15:40 +0200 Subject: [PATCH 1/4] Update storage of cache cheatnet state with the txs state --- crates/cheatnet/src/forking/cache.rs | 6 +- crates/cheatnet/src/forking/state.rs | 119 ++++++++++++++++++++++++--- 2 files changed, 113 insertions(+), 12 deletions(-) diff --git a/crates/cheatnet/src/forking/cache.rs b/crates/cheatnet/src/forking/cache.rs index c5222dbf26..218eb94a23 100644 --- a/crates/cheatnet/src/forking/cache.rs +++ b/crates/cheatnet/src/forking/cache.rs @@ -81,7 +81,7 @@ impl ToString for ForkCacheContent { } } -#[derive(Debug)] +#[derive(Debug, Default)] pub struct ForkCache { fork_cache_content: ForkCacheContent, cache_file: Utf8PathBuf, @@ -89,7 +89,9 @@ pub struct ForkCache { impl Drop for ForkCache { fn drop(&mut self) { - self.save(); + if !self.cache_file.as_str().is_empty() { + self.save(); + } } } diff --git a/crates/cheatnet/src/forking/state.rs b/crates/cheatnet/src/forking/state.rs index c818344d52..dc5b6384ea 100644 --- a/crates/cheatnet/src/forking/state.rs +++ b/crates/cheatnet/src/forking/state.rs @@ -14,8 +14,8 @@ use flate2::read::GzDecoder; use num_bigint::BigUint; use runtime::starknet::context::SerializableGasPrices; use starknet::core::types::{ - BlockId, ContractClass as ContractClassStarknet, FieldElement, MaybePendingBlockWithTxHashes, - StarknetError, + BlockId, ContractClass as ContractClassStarknet, ContractStorageDiffItem, FieldElement, + MaybePendingBlockWithTxHashes, StarknetError, TransactionTrace, }; use starknet::providers::jsonrpc::HttpTransport; use starknet::providers::{JsonRpcClient, Provider, ProviderError}; @@ -40,15 +40,34 @@ pub struct ForkStateReader { } impl ForkStateReader { - pub fn new(url: Url, block_number: BlockNumber, cache_dir: &str) -> Result { - Ok(ForkStateReader { - cache: RefCell::new( - ForkCache::load_or_new(&url, block_number, cache_dir) - .context("Could not create fork cache")?, - ), - client: JsonRpcClient::new(HttpTransport::new(url)), + pub fn new( + url: Url, + block_number: BlockNumber, + transaction_index: usize, + cache_dir: &str, + ) -> Result { + let mut fork_cache = ForkCache::default(); + let mut fork_state_reader = ForkStateReader { + cache: RefCell::new(fork_cache), + client: JsonRpcClient::new(HttpTransport::new(url.clone())), block_number, - }) + }; + + //Get over all transaction till transaction_index and store new storage values in + //fork_cache_content + fork_state_reader.get_transactions_storage_diff( + BlockId::Number(block_number.0 + 1), + transaction_index, + )?; + + //load new fork_cache with updated fork_cache_content + fork_cache = ForkCache::load_or_new(&url, block_number, cache_dir) + .context("Could not load fork cache")?; + + fork_state_reader.cache = RefCell::new(fork_cache); + + // Return the initialized and state updated ForkStateReader + Ok(fork_state_reader) } fn block_id(&self) -> BlockId { @@ -64,6 +83,86 @@ impl ForkStateReader { .get_compiled_contract_class(&class_hash) .cloned() } + + pub fn get_transactions_storage_diff( + &self, + block_id: BlockId, + transaction_index: usize, + ) -> Result<(), StateError> { + let results = tokio::task::block_in_place(|| { + tokio::runtime::Handle::current() + .block_on(self.client.trace_block_transactions(block_id)) + }) + .map_err(|_| { + StateError::StateReadError("Failed to trace block transactions".to_string()) + })?; + + for (index, result) in results.into_iter().enumerate() { + if index != transaction_index { + match result.trace_root { + TransactionTrace::Invoke(invoke_trace) => { + if let Some(state_diff) = invoke_trace.state_diff { + let contract_storage_diff = state_diff.storage_diffs; + // Process the storage differences and update cache storage content + self.process_storage_diffs(&contract_storage_diff)?; + } + return Ok(()); + } + TransactionTrace::Declare(_) => { + //TODO implement for Declare tx + return Ok(()); + } + TransactionTrace::DeployAccount(_) => { + return Err(StateError::StateReadError( + "DeployAccount account trace not found".to_string(), + )); + } + TransactionTrace::L1Handler(_) => { + return Err(StateError::StateReadError( + "L1Handler trace not found".to_string(), + )); + } + } + } + } + + Err(StateError::StateReadError( + "Transactions trace not found".to_string(), + )) + } + + fn process_storage_diffs( + &self, + storage_diffs: &[ContractStorageDiffItem], + ) -> Result<(), StateError> { + for storage_diff in storage_diffs.iter() { + let contract_address = &storage_diff.address; + + for storage_entry in storage_diff.storage_entries.iter() { + let key = StorageKey::try_from(StarkFelt::from(storage_entry.key))?; + let new_value: StarkFelt = storage_entry.value.into_(); + let contract_address: ContractAddress = + ContractAddress::try_from(StarkFelt::from(*contract_address))?; + + match self.get_storage_at(contract_address, key) { + Ok(current_value) => { + if current_value != new_value { + self.cache.borrow_mut().cache_get_storage_at( + contract_address, + key, + new_value, + ); + } + } + Err(e) => { + return Err(e); + } + } + } + } + + Ok(()) + } } #[allow(clippy::needless_pass_by_value)] From 62d43503f6acdb596bba93dc376717b241bb95bf Mon Sep 17 00:00:00 2001 From: marijamijailovic Date: Tue, 27 Aug 2024 19:14:59 +0200 Subject: [PATCH 2/4] Store storage diff state for all txs till the one we are simulate --- crates/cheatnet/src/forking/cache.rs | 4 +- crates/cheatnet/src/forking/state.rs | 113 +++++++++++++-------------- 2 files changed, 55 insertions(+), 62 deletions(-) diff --git a/crates/cheatnet/src/forking/cache.rs b/crates/cheatnet/src/forking/cache.rs index 218eb94a23..a15a84a4ca 100644 --- a/crates/cheatnet/src/forking/cache.rs +++ b/crates/cheatnet/src/forking/cache.rs @@ -89,9 +89,7 @@ pub struct ForkCache { impl Drop for ForkCache { fn drop(&mut self) { - if !self.cache_file.as_str().is_empty() { - self.save(); - } + self.save(); } } diff --git a/crates/cheatnet/src/forking/state.rs b/crates/cheatnet/src/forking/state.rs index dc5b6384ea..0cada62436 100644 --- a/crates/cheatnet/src/forking/state.rs +++ b/crates/cheatnet/src/forking/state.rs @@ -37,6 +37,7 @@ pub struct ForkStateReader { client: JsonRpcClient, block_number: BlockNumber, cache: RefCell, + storage_diff: HashMap>, } impl ForkStateReader { @@ -46,26 +47,22 @@ impl ForkStateReader { transaction_index: usize, cache_dir: &str, ) -> Result { - let mut fork_cache = ForkCache::default(); + let mut fork_cache = ForkCache::load_or_new(&url, block_number, cache_dir) + .context("Could not load fork cache")?; let mut fork_state_reader = ForkStateReader { cache: RefCell::new(fork_cache), client: JsonRpcClient::new(HttpTransport::new(url.clone())), - block_number, + block_number: BlockNumber(block_number.0), + storage_diff: HashMap::new(), }; //Get over all transaction till transaction_index and store new storage values in - //fork_cache_content + //storage_diff hash map fork_state_reader.get_transactions_storage_diff( BlockId::Number(block_number.0 + 1), transaction_index, )?; - //load new fork_cache with updated fork_cache_content - fork_cache = ForkCache::load_or_new(&url, block_number, cache_dir) - .context("Could not load fork cache")?; - - fork_state_reader.cache = RefCell::new(fork_cache); - // Return the initialized and state updated ForkStateReader Ok(fork_state_reader) } @@ -85,7 +82,7 @@ impl ForkStateReader { } pub fn get_transactions_storage_diff( - &self, + &mut self, block_id: BlockId, transaction_index: usize, ) -> Result<(), StateError> { @@ -98,70 +95,54 @@ impl ForkStateReader { })?; for (index, result) in results.into_iter().enumerate() { - if index != transaction_index { - match result.trace_root { - TransactionTrace::Invoke(invoke_trace) => { - if let Some(state_diff) = invoke_trace.state_diff { - let contract_storage_diff = state_diff.storage_diffs; - // Process the storage differences and update cache storage content - self.process_storage_diffs(&contract_storage_diff)?; - } - return Ok(()); + if index == transaction_index { + break; + } + match &result.trace_root { + TransactionTrace::Invoke(invoke_trace) => { + if let Some(state_diff) = &invoke_trace.state_diff { + let contract_storage_diff = &state_diff.storage_diffs; + self.collect_storage_diffs(contract_storage_diff); } - TransactionTrace::Declare(_) => { - //TODO implement for Declare tx - return Ok(()); + } + TransactionTrace::Declare(declare_trace) => { + if let Some(state_diff) = &declare_trace.state_diff { + let contract_storage_diff = &state_diff.storage_diffs; + self.collect_storage_diffs(contract_storage_diff); } - TransactionTrace::DeployAccount(_) => { - return Err(StateError::StateReadError( - "DeployAccount account trace not found".to_string(), - )); + } + TransactionTrace::DeployAccount(deploy_trace) => { + if let Some(state_diff) = &deploy_trace.state_diff { + let contract_storage_diff = &state_diff.storage_diffs; + self.collect_storage_diffs(contract_storage_diff); } - TransactionTrace::L1Handler(_) => { - return Err(StateError::StateReadError( - "L1Handler trace not found".to_string(), - )); + } + TransactionTrace::L1Handler(l1handler_trace) => { + if let Some(state_diff) = &l1handler_trace.state_diff { + let contract_storage_diff = &state_diff.storage_diffs; + self.collect_storage_diffs(contract_storage_diff); } } } } - Err(StateError::StateReadError( - "Transactions trace not found".to_string(), - )) + Ok(()) } - fn process_storage_diffs( - &self, - storage_diffs: &[ContractStorageDiffItem], - ) -> Result<(), StateError> { + fn collect_storage_diffs(&mut self, storage_diffs: &[ContractStorageDiffItem]) { for storage_diff in storage_diffs.iter() { - let contract_address = &storage_diff.address; - + let contract_address: ContractAddress = + ContractAddress::try_from(StarkFelt::from(storage_diff.address)).unwrap(); + let contract_storage = self.storage_diff.entry(contract_address).or_default(); for storage_entry in storage_diff.storage_entries.iter() { - let key = StorageKey::try_from(StarkFelt::from(storage_entry.key))?; + let key = StorageKey::try_from(StarkFelt::from(storage_entry.key)).unwrap(); let new_value: StarkFelt = storage_entry.value.into_(); - let contract_address: ContractAddress = - ContractAddress::try_from(StarkFelt::from(*contract_address))?; - - match self.get_storage_at(contract_address, key) { - Ok(current_value) => { - if current_value != new_value { - self.cache.borrow_mut().cache_get_storage_at( - contract_address, - key, - new_value, - ); - } - } - Err(e) => { - return Err(e); - } - } + contract_storage.insert(key, new_value); + // self.cache + // .borrow_mut() + // .cache_get_storage_at(contract_address, key, new_value); } } - - Ok(()) } } @@ -220,10 +201,23 @@ impl StateReader for ForkStateReader { contract_address: ContractAddress, key: StorageKey, ) -> StateResult { + // First check cache if let Some(cache_hit) = self.cache.borrow().get_storage_at(&contract_address, &key) { return Ok(cache_hit); } + // Second check the storage_diff hash map + if let Some(contract_updates) = self.storage_diff.get(&contract_address) { + if let Some(&value) = contract_updates.get(&key) { + self.cache + .borrow_mut() + .cache_get_storage_at(contract_address, key, value); + + return Ok(value); + } + } + + // Thrird ping provider match tokio::task::block_in_place(|| { tokio::runtime::Handle::current().block_on(self.client.get_storage_at( FieldElement::from_(contract_address), @@ -231,6 +225,7 @@ impl StateReader for ForkStateReader { self.block_id(), )) }) { + // match self.runtime.block_on(self.client.get_storage_at( // FieldElement::from_(contract_address), // FieldElement::from_(*key.0.key()), From 367427cca6253a8a4de6fac795b212f60403554b Mon Sep 17 00:00:00 2001 From: marijamijailovic Date: Thu, 29 Aug 2024 13:19:25 +0200 Subject: [PATCH 3/4] Remove comment, typo fix --- crates/cheatnet/src/forking/state.rs | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/crates/cheatnet/src/forking/state.rs b/crates/cheatnet/src/forking/state.rs index 0cada62436..50c80df938 100644 --- a/crates/cheatnet/src/forking/state.rs +++ b/crates/cheatnet/src/forking/state.rs @@ -138,9 +138,6 @@ impl ForkStateReader { let key = StorageKey::try_from(StarkFelt::from(storage_entry.key)).unwrap(); let new_value: StarkFelt = storage_entry.value.into_(); contract_storage.insert(key, new_value); - // self.cache - // .borrow_mut() - // .cache_get_storage_at(contract_address, key, new_value); } } } @@ -217,7 +214,7 @@ impl StateReader for ForkStateReader { } } - // Thrird ping provider + // Third ping provider match tokio::task::block_in_place(|| { tokio::runtime::Handle::current().block_on(self.client.get_storage_at( FieldElement::from_(contract_address), From c61491c3803dca2881fbfbe8178ff5758626f789 Mon Sep 17 00:00:00 2001 From: marijamijailovic Date: Fri, 30 Aug 2024 13:51:30 +0200 Subject: [PATCH 4/4] Added check to skip storage update if we have 1 tx --- crates/cheatnet/src/forking/state.rs | 40 +++++++++++++++++++++------- 1 file changed, 30 insertions(+), 10 deletions(-) diff --git a/crates/cheatnet/src/forking/state.rs b/crates/cheatnet/src/forking/state.rs index 50c80df938..18e6003878 100644 --- a/crates/cheatnet/src/forking/state.rs +++ b/crates/cheatnet/src/forking/state.rs @@ -47,7 +47,7 @@ impl ForkStateReader { transaction_index: usize, cache_dir: &str, ) -> Result { - let mut fork_cache = ForkCache::load_or_new(&url, block_number, cache_dir) + let fork_cache = ForkCache::load_or_new(&url, block_number, cache_dir) .context("Could not load fork cache")?; let mut fork_state_reader = ForkStateReader { cache: RefCell::new(fork_cache), @@ -56,13 +56,17 @@ impl ForkStateReader { storage_diff: HashMap::new(), }; - //Get over all transaction till transaction_index and store new storage values in - //storage_diff hash map - fork_state_reader.get_transactions_storage_diff( - BlockId::Number(block_number.0 + 1), - transaction_index, - )?; - + let real_block_id = BlockId::Number(block_number.0 + 1); + let tx_in_block = fork_state_reader + .get_block_transaction_count(real_block_id) + .context("Unable to get block transactions count from node provider")?; + if tx_in_block > 1 { + //Get over all transaction till transaction_index and store new storage values in + //storage_diff hash map + fork_state_reader + .get_transactions_storage_diff(real_block_id, transaction_index) + .context("Unable to get trace block transactions from node provider")?; + } // Return the initialized and state updated ForkStateReader Ok(fork_state_reader) } @@ -81,6 +85,20 @@ impl ForkStateReader { .cloned() } + pub fn get_block_transaction_count(&self, block_id: BlockId) -> Result { + let result = tokio::task::block_in_place(|| { + tokio::runtime::Handle::current() + .block_on(self.client.get_block_transaction_count(block_id)) + }) + .map_err(|err| { + StateError::StateReadError(format!( + "Unable to get block transactions count from fork ({err})" + )) + })?; + + Ok(result) + } + pub fn get_transactions_storage_diff( &mut self, block_id: BlockId, @@ -90,8 +108,10 @@ impl ForkStateReader { tokio::runtime::Handle::current() .block_on(self.client.trace_block_transactions(block_id)) }) - .map_err(|_| { - StateError::StateReadError("Failed to trace block transactions".to_string()) + .map_err(|err| { + StateError::StateReadError(format!( + "Unable to get trace block transactions from fork ({err})" + )) })?; for (index, result) in results.into_iter().enumerate() {