diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index 6fd6ee24e..aad1f7dfd 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -1244,6 +1244,21 @@ impl Cheatcodes { } } + // record immutable variables + if result.execution_result.is_success() { + for (addr, imm_values) in result.recorded_immutables { + let addr = addr.to_address(); + let keys = imm_values + .into_keys() + .map(|slot_index| { + foundry_zksync_core::get_immutable_slot_key(addr, slot_index) + .to_ru256() + }) + .collect::>(); + ecx.db.save_zk_immutable_storage(addr, keys); + } + } + match result.execution_result { ExecutionResult::Success { output, gas_used, .. } => { let _ = gas.record_cost(gas_used); diff --git a/crates/evm/core/src/backend/cow.rs b/crates/evm/core/src/backend/cow.rs index e1879c33a..c661729b0 100644 --- a/crates/evm/core/src/backend/cow.rs +++ b/crates/evm/core/src/backend/cow.rs @@ -21,7 +21,10 @@ use revm::{ }, Database, DatabaseCommit, JournaledState, }; -use std::{borrow::Cow, collections::BTreeMap}; +use std::{ + borrow::Cow, + collections::{BTreeMap, HashSet}, +}; /// A wrapper around `Backend` that ensures only `revm::DatabaseRef` functions are called. /// @@ -133,6 +136,10 @@ impl DatabaseExt for CowBackend<'_> { self.backend.to_mut().get_fork_info(id) } + fn save_zk_immutable_storage(&mut self, addr: Address, keys: HashSet) { + self.backend.to_mut().save_zk_immutable_storage(addr, keys) + } + fn snapshot_state(&mut self, journaled_state: &JournaledState, env: &Env) -> U256 { self.backend_mut(env).snapshot_state(journaled_state, env) } diff --git a/crates/evm/core/src/backend/mod.rs b/crates/evm/core/src/backend/mod.rs index c6ed50ddb..fed2ecba6 100644 --- a/crates/evm/core/src/backend/mod.rs +++ b/crates/evm/core/src/backend/mod.rs @@ -15,7 +15,8 @@ use eyre::Context; use foundry_common::{is_known_system_sender, SYSTEM_TRANSACTION_TYPE}; pub use foundry_fork_db::{cache::BlockchainDbMeta, BlockchainDb, SharedBackend}; use foundry_zksync_core::{ - convert::ConvertH160, ACCOUNT_CODE_STORAGE_ADDRESS, L2_BASE_TOKEN_ADDRESS, NONCE_HOLDER_ADDRESS, + convert::ConvertH160, ACCOUNT_CODE_STORAGE_ADDRESS, IMMUTABLE_SIMULATOR_STORAGE_ADDRESS, + KNOWN_CODES_STORAGE_ADDRESS, L2_BASE_TOKEN_ADDRESS, NONCE_HOLDER_ADDRESS, }; use itertools::Itertools; use revm::{ @@ -100,6 +101,14 @@ pub trait DatabaseExt: Database + DatabaseCommit { /// and the the fork environment. fn get_fork_info(&mut self, id: LocalForkId) -> eyre::Result; + /// Saves the storage keys for immutable variables per address. + /// + /// These are required during fork to help merge the persisted addresses, as they are stored + /// hashed so there is currently no way to retrieve all the address associated storage keys. + /// We store all the storage keys here, even if the addresses are not marked persistent as + /// they can be marked at a later stage as well. + fn save_zk_immutable_storage(&mut self, addr: Address, keys: HashSet); + /// Reverts the snapshot if it exists /// /// Returns `true` if the snapshot was successfully reverted, `false` if no snapshot for that id @@ -491,6 +500,8 @@ pub struct Backend { /// The balance, nonce and code are stored under zkSync's respective system contract /// storages. These need to be merged into the forked storage. pub is_zk: bool, + /// Store storage keys per contract address for immutable variables. + zk_recorded_immutable_keys: HashMap>, } impl Backend { @@ -524,6 +535,7 @@ impl Backend { inner, fork_url_type: Default::default(), is_zk: false, + zk_recorded_immutable_keys: Default::default(), }; if let Some(fork) = fork { @@ -564,6 +576,7 @@ impl Backend { inner: Default::default(), fork_url_type: Default::default(), is_zk: false, + zk_recorded_immutable_keys: Default::default(), } } @@ -670,13 +683,13 @@ impl Backend { &self, active_journaled_state: &mut JournaledState, target_fork: &mut Fork, - merge_zk_db: bool, + zk_state: Option, ) { self.update_fork_db_contracts( self.inner.persistent_accounts.iter().copied(), active_journaled_state, target_fork, - merge_zk_db, + zk_state, ) } @@ -686,17 +699,17 @@ impl Backend { accounts: impl IntoIterator, active_journaled_state: &mut JournaledState, target_fork: &mut Fork, - merge_zk_db: bool, + zk_state: Option, ) { if let Some(db) = self.active_fork_db() { - merge_account_data(accounts, db, active_journaled_state, target_fork, merge_zk_db) + merge_account_data(accounts, db, active_journaled_state, target_fork, zk_state) } else { merge_account_data( accounts, &self.mem_db, active_journaled_state, target_fork, - merge_zk_db, + zk_state, ) } } @@ -984,6 +997,13 @@ impl DatabaseExt for Backend { Ok(ForkInfo { fork_type, fork_env }) } + fn save_zk_immutable_storage(&mut self, addr: Address, keys: HashSet) { + self.zk_recorded_immutable_keys + .entry(addr) + .and_modify(|entry| entry.extend(&keys)) + .or_insert(keys); + } + fn snapshot_state(&mut self, journaled_state: &JournaledState, env: &Env) -> U256 { trace!("create snapshot"); let id = self.inner.state_snapshots.insert(BackendStateSnapshot::new( @@ -1144,6 +1164,9 @@ impl DatabaseExt for Backend { .map(|url| self.fork_url_type.get(&url).is_zk()) .unwrap_or_default(); let merge_zk_db = is_current_zk_fork && is_target_zk_fork; + let zk_state = merge_zk_db.then(|| ZkMergeState { + persistent_immutable_keys: self.zk_recorded_immutable_keys.clone(), + }); let fork_env = self .forks @@ -1212,7 +1235,7 @@ impl DatabaseExt for Backend { caller_account.into() }); - self.update_fork_db(active_journaled_state, &mut fork, merge_zk_db); + self.update_fork_db(active_journaled_state, &mut fork, zk_state); // insert the fork back self.inner.set_fork(idx, fork); @@ -1930,6 +1953,12 @@ pub(crate) fn update_current_env_with_fork_env(current: &mut Env, fork: Env) { current.tx.chain_id = fork.tx.chain_id; } +/// Defines the zksync specific state to help during merge. +#[derive(Debug, Default)] +pub(crate) struct ZkMergeState { + persistent_immutable_keys: HashMap>, +} + /// Clones the data of the given `accounts` from the `active` database into the `fork_db` /// This includes the data held in storage (`CacheDB`) and kept in the `JournaledState`. pub(crate) fn merge_account_data( @@ -1937,19 +1966,20 @@ pub(crate) fn merge_account_data( active: &CacheDB, active_journaled_state: &mut JournaledState, target_fork: &mut Fork, - merge_zk_db: bool, + zk_state: Option, ) { for addr in accounts.into_iter() { merge_db_account_data(addr, active, &mut target_fork.db); - if merge_zk_db { - merge_zk_account_data(addr, active, &mut target_fork.db); + if let Some(zk_state) = &zk_state { + merge_zk_account_data(addr, active, &mut target_fork.db, zk_state); } merge_journaled_state_data(addr, active_journaled_state, &mut target_fork.journaled_state); - if merge_zk_db { + if let Some(zk_state) = &zk_state { merge_zk_journaled_state_data( addr, active_journaled_state, &mut target_fork.journaled_state, + zk_state, ); } } @@ -2020,53 +2050,66 @@ fn merge_zk_account_data( addr: Address, active: &CacheDB, fork_db: &mut ForkDB, + _zk_state: &ZkMergeState, ) { - let mut merge_system_contract_entry = |system_contract: Address, slot: U256| { - let Some(acc) = active.accounts.get(&system_contract) else { return }; - - // port contract cache over - if let Some(code) = active.contracts.get(&acc.info.code_hash) { - trace!("merging contract cache"); - fork_db.contracts.insert(acc.info.code_hash, code.clone()); - } - - // prepare only the specified slot in account storage - let mut new_acc = acc.clone(); - new_acc.storage = Default::default(); - if let Some(value) = acc.storage.get(&slot) { - new_acc.storage.insert(slot, *value); - } + let merge_system_contract_entry = + |fork_db: &mut ForkDB, system_contract: Address, slot: U256| { + let Some(acc) = active.accounts.get(&system_contract) else { return }; + + // port contract cache over + if let Some(code) = active.contracts.get(&acc.info.code_hash) { + trace!("merging contract cache"); + fork_db.contracts.insert(acc.info.code_hash, code.clone()); + } - // port account storage over - match fork_db.accounts.entry(system_contract) { - Entry::Vacant(vacant) => { - trace!("target account not present - inserting from active"); - // if the fork_db doesn't have the target account - // insert the entire thing - vacant.insert(new_acc); + // prepare only the specified slot in account storage + let mut new_acc = acc.clone(); + new_acc.storage = Default::default(); + if let Some(value) = acc.storage.get(&slot) { + new_acc.storage.insert(slot, *value); } - Entry::Occupied(mut occupied) => { - trace!("target account present - merging storage slots"); - // if the fork_db does have the system, - // extend the existing storage (overriding) - let fork_account = occupied.get_mut(); - fork_account.storage.extend(&new_acc.storage); + + // port account storage over + match fork_db.accounts.entry(system_contract) { + Entry::Vacant(vacant) => { + trace!("target account not present - inserting from active"); + // if the fork_db doesn't have the target account + // insert the entire thing + vacant.insert(new_acc); + } + Entry::Occupied(mut occupied) => { + trace!("target account present - merging storage slots"); + // if the fork_db does have the system, + // extend the existing storage (overriding) + let fork_account = occupied.get_mut(); + fork_account.storage.extend(&new_acc.storage); + } } - } - }; + }; merge_system_contract_entry( + fork_db, L2_BASE_TOKEN_ADDRESS.to_address(), foundry_zksync_core::get_balance_key(addr), ); merge_system_contract_entry( + fork_db, ACCOUNT_CODE_STORAGE_ADDRESS.to_address(), foundry_zksync_core::get_account_code_key(addr), ); merge_system_contract_entry( + fork_db, NONCE_HOLDER_ADDRESS.to_address(), foundry_zksync_core::get_nonce_key(addr), ); + + if let Some(acc) = active.accounts.get(&addr) { + merge_system_contract_entry( + fork_db, + KNOWN_CODES_STORAGE_ADDRESS.to_address(), + U256::from_be_slice(&acc.info.code_hash.0[..]), + ); + } } /// Clones the account data from the `active_journaled_state` into the `fork_journaled_state` for @@ -2075,40 +2118,61 @@ fn merge_zk_journaled_state_data( addr: Address, active_journaled_state: &JournaledState, fork_journaled_state: &mut JournaledState, + zk_state: &ZkMergeState, ) { - let mut merge_system_contract_entry = |system_contract: Address, slot: U256| { - if let Some(acc) = active_journaled_state.state.get(&system_contract) { - // prepare only the specified slot in account storage - let mut new_acc = acc.clone(); - new_acc.storage = Default::default(); - if let Some(value) = acc.storage.get(&slot).cloned() { - new_acc.storage.insert(slot, value); - } - - match fork_journaled_state.state.entry(system_contract) { - Entry::Vacant(vacant) => { - vacant.insert(new_acc); + let merge_system_contract_entry = + |fork_journaled_state: &mut JournaledState, system_contract: Address, slot: U256| { + if let Some(acc) = active_journaled_state.state.get(&system_contract) { + // prepare only the specified slot in account storage + let mut new_acc = acc.clone(); + new_acc.storage = Default::default(); + if let Some(value) = acc.storage.get(&slot).cloned() { + new_acc.storage.insert(slot, value); } - Entry::Occupied(mut occupied) => { - let fork_account = occupied.get_mut(); - fork_account.storage.extend(new_acc.storage); + + match fork_journaled_state.state.entry(system_contract) { + Entry::Vacant(vacant) => { + vacant.insert(new_acc); + } + Entry::Occupied(mut occupied) => { + let fork_account = occupied.get_mut(); + fork_account.storage.extend(new_acc.storage); + } } } - } - }; + }; merge_system_contract_entry( + fork_journaled_state, L2_BASE_TOKEN_ADDRESS.to_address(), foundry_zksync_core::get_balance_key(addr), ); merge_system_contract_entry( + fork_journaled_state, ACCOUNT_CODE_STORAGE_ADDRESS.to_address(), foundry_zksync_core::get_account_code_key(addr), ); merge_system_contract_entry( + fork_journaled_state, NONCE_HOLDER_ADDRESS.to_address(), foundry_zksync_core::get_nonce_key(addr), ); + + if let Some(acc) = active_journaled_state.state.get(&addr) { + merge_system_contract_entry( + fork_journaled_state, + KNOWN_CODES_STORAGE_ADDRESS.to_address(), + U256::from_be_slice(&acc.info.code_hash.0[..]), + ); + } + + // merge immutable storage. + let immutable_simulator_addr = IMMUTABLE_SIMULATOR_STORAGE_ADDRESS.to_address(); + if let Some(immutable_storage_keys) = zk_state.persistent_immutable_keys.get(&addr) { + for slot_key in immutable_storage_keys { + merge_system_contract_entry(fork_journaled_state, immutable_simulator_addr, *slot_key); + } + } } /// Returns true of the address is a contract diff --git a/crates/forge/tests/it/zk/fork.rs b/crates/forge/tests/it/zk/fork.rs index 5fed33832..e8477753e 100644 --- a/crates/forge/tests/it/zk/fork.rs +++ b/crates/forge/tests/it/zk/fork.rs @@ -12,3 +12,11 @@ async fn test_zk_setup_fork_failure() { TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await; } + +#[tokio::test(flavor = "multi_thread")] +async fn test_zk_immutable_vars_persist_after_fork() { + let runner = TEST_DATA_DEFAULT.runner_zksync(); + let filter = Filter::new(".*", "ZkForkImmutableVarsTest", ".*"); + + TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await; +} diff --git a/crates/zksync/core/src/lib.rs b/crates/zksync/core/src/lib.rs index e2f63be4d..80af94fcd 100644 --- a/crates/zksync/core/src/lib.rs +++ b/crates/zksync/core/src/lib.rs @@ -19,7 +19,7 @@ pub mod vm; pub mod state; use alloy_network::{AnyNetwork, TxSigner}; -use alloy_primitives::{address, hex, Address, Bytes, U256 as rU256}; +use alloy_primitives::{address, hex, keccak256, Address, Bytes, U256 as rU256}; use alloy_provider::Provider; use alloy_rpc_types::TransactionRequest; use alloy_serde::WithOtherFields; @@ -37,11 +37,12 @@ pub use vm::{balance, encode_create_params, nonce}; pub use vm::{SELECTOR_CONTRACT_DEPLOYER_CREATE, SELECTOR_CONTRACT_DEPLOYER_CREATE2}; pub use zksync_multivm::interface::{Call, CallType}; -use zksync_types::utils::storage_key_for_eth_balance; pub use zksync_types::{ - ACCOUNT_CODE_STORAGE_ADDRESS, CONTRACT_DEPLOYER_ADDRESS, H256, L2_BASE_TOKEN_ADDRESS, + ethabi, ACCOUNT_CODE_STORAGE_ADDRESS, CONTRACT_DEPLOYER_ADDRESS, H256, + IMMUTABLE_SIMULATOR_STORAGE_ADDRESS, KNOWN_CODES_STORAGE_ADDRESS, L2_BASE_TOKEN_ADDRESS, NONCE_HOLDER_ADDRESS, }; +use zksync_types::{utils::storage_key_for_eth_balance, U256}; pub use zksync_utils::bytecode::hash_bytecode; use zksync_web3_rs::{ eip712::{Eip712Meta, Eip712Transaction, Eip712TransactionRequest, PaymasterParams}, @@ -303,3 +304,41 @@ pub fn try_decode_create(data: &[u8]) -> Result<(H256, Vec)> { Ok((H256(bytecode_hash.0), constructor_args.to_vec())) } + +/// Gets the mapping key for the `ImmutableSimulator::immutableDataStorage`. +/// +/// This retrieves the key for a given contract address and variable slot. +/// See https://github.com/matter-labs/era-contracts/blob/main/system-contracts/contracts/ImmutableSimulator.sol#L21 +pub fn get_immutable_slot_key(address: Address, slot_index: rU256) -> H256 { + let immutable_data_storage_key = keccak256(ethabi::encode(&[ + ethabi::Token::Address(address.to_h160()), + ethabi::Token::Uint(U256::zero()), + ])); + let immutable_data_storage_key = H256(*immutable_data_storage_key); + + let immutable_value_key = keccak256(ethabi::encode(&[ + ethabi::Token::Uint(slot_index.to_u256()), + ethabi::Token::FixedBytes(immutable_data_storage_key.to_fixed_bytes().to_vec()), + ])); + + H256(*immutable_value_key) +} + +#[cfg(test)] +mod tests { + use std::str::FromStr; + + use super::*; + + #[test] + fn test_get_immutable_slot_key() { + let actual_key = get_immutable_slot_key( + address!("f9e9ba9ed9b96ab918c74b21dd0f1d5f2ac38a30"), + rU256::from(10u32), + ); + let expected_key = + H256::from_str("db259b642223206a098c9ffaaf8e4bfd2d60060e8365bb349b2ea2b720d9837c") + .expect("invalid h256"); + assert_eq!(expected_key, actual_key) + } +} diff --git a/crates/zksync/core/src/vm/inspect.rs b/crates/zksync/core/src/vm/inspect.rs index 5742fba60..96f22f24a 100644 --- a/crates/zksync/core/src/vm/inspect.rs +++ b/crates/zksync/core/src/vm/inspect.rs @@ -1,4 +1,4 @@ -use alloy_primitives::{hex, Log}; +use alloy_primitives::{hex, FixedBytes, Log}; use era_test_node::{ config::node::ShowCalls, formatter, @@ -68,6 +68,8 @@ pub struct ZKVMExecutionResult { pub execution_result: rExecutionResult, /// Call traces pub call_traces: Vec, + /// Immutables recorded via calls to ImmutableSimulator::setImmutables. + pub recorded_immutables: rHashMap>>, } /// Revm-style result with ZKVM Execution @@ -112,6 +114,7 @@ where logs: result.logs, call_traces: result.call_traces, execution_result: exec, + recorded_immutables: result.recorded_immutables, }); } (None, exec) => { @@ -119,6 +122,7 @@ where logs: result.logs, call_traces: result.call_traces, execution_result: exec, + recorded_immutables: result.recorded_immutables, }); } ( @@ -133,11 +137,13 @@ where logs: agg_logs, output: agg_output, }, + recorded_immutables: aggregated_recorded_immutables, }), rExecutionResult::Success { reason, gas_used, gas_refunded, logs, output }, ) => { aggregated_logs.append(&mut result.logs); aggregated_call_traces.append(&mut result.call_traces); + aggregated_recorded_immutables.extend(result.recorded_immutables); *agg_reason = reason; *agg_gas_used += gas_used; *agg_gas_refunded += gas_refunded; @@ -199,6 +205,7 @@ where bytecodes, modified_storage, call_traces, + recorded_immutables, create_outcome, gas_usage, } = inspect_inner(tx, storage_ptr, chain_id, ccx, call_ctx); @@ -270,6 +277,7 @@ where logs, output, }, + recorded_immutables, } } ExecutionResult::Revert { output } => { @@ -286,6 +294,7 @@ where gas_used: gas_usage.gas_used(), output: Bytes::from(output), }, + recorded_immutables, } } ExecutionResult::Halt { reason } => { @@ -302,6 +311,7 @@ where reason: mapped_reason, gas_used: gas_usage.gas_used(), }, + recorded_immutables, } } }; @@ -410,6 +420,7 @@ struct InnerZkVmResult { call_traces: Vec, create_outcome: Option, gas_usage: ZkVmGasUsage, + recorded_immutables: rHashMap>>, } fn inspect_inner( @@ -478,6 +489,7 @@ fn inspect_inner( if let Some(expected_calls) = ccx.expected_calls.as_mut() { expected_calls.extend(cheatcode_result.expected_calls); } + let recorded_immutables = cheatcode_result.recorded_immutables; // populate gas usage info let bootloader_debug = Arc::try_unwrap(bootloader_debug_tracer_result) @@ -548,11 +560,7 @@ fn inspect_inner( .expect("failed converting bytecode to factory dep") }) .collect::>>(); - let modified_storage = if is_static { - Default::default() - } else { - storage.borrow().modified_storage_keys().clone() - }; + let modified_storage = storage.borrow().modified_storage_keys().clone(); // patch CREATE traces. for call in call_traces.iter_mut() { @@ -589,13 +597,26 @@ fn inspect_inner( None }; - InnerZkVmResult { - tx_result, - bytecodes, - modified_storage, - call_traces, - create_outcome, - gas_usage, + if is_static { + InnerZkVmResult { + tx_result, + bytecodes: Default::default(), + modified_storage: Default::default(), + call_traces, + create_outcome, + gas_usage, + recorded_immutables: Default::default(), + } + } else { + InnerZkVmResult { + tx_result, + bytecodes, + modified_storage, + call_traces, + create_outcome, + gas_usage, + recorded_immutables, + } } } diff --git a/crates/zksync/core/src/vm/tracers/cheatcode.rs b/crates/zksync/core/src/vm/tracers/cheatcode.rs index b738adabb..3f5f879bd 100644 --- a/crates/zksync/core/src/vm/tracers/cheatcode.rs +++ b/crates/zksync/core/src/vm/tracers/cheatcode.rs @@ -4,7 +4,7 @@ use std::{ sync::Arc, }; -use alloy_primitives::{hex, map::HashMap, Address, Bytes, U256 as rU256}; +use alloy_primitives::{hex, map::HashMap, Address, Bytes, FixedBytes, U256 as rU256}; use foundry_cheatcodes_common::{ expect::ExpectedCallTracker, mock::{MockCallDataContext, MockCallReturnData}, @@ -21,8 +21,8 @@ use zksync_multivm::{ }; use zksync_state::interface::{ReadStorage, StoragePtr, WriteStorage}; use zksync_types::{ - get_code_key, StorageValue, BOOTLOADER_ADDRESS, CONTRACT_DEPLOYER_ADDRESS, H256, - SYSTEM_CONTEXT_ADDRESS, U256, + ethabi, get_code_key, StorageValue, BOOTLOADER_ADDRESS, CONTRACT_DEPLOYER_ADDRESS, H160, H256, + IMMUTABLE_SIMULATOR_STORAGE_ADDRESS, SYSTEM_CONTEXT_ADDRESS, U256, }; use zksync_utils::bytecode::hash_bytecode; @@ -70,6 +70,13 @@ const SELECTOR_BASE_FEE: [u8; 4] = hex!("6ef25c3a"); /// Selector for `getBlockHashEVM(uint256)` const SELECTOR_BLOCK_HASH: [u8; 4] = hex!("80b41246"); +/// Selector for setting immutables for an address. +/// This is used to retrieve the immutables and use them in merging storage +/// during forks. +/// +/// Selector for `setImmutables(address, (uint256,bytes32)[])", +const SELECTOR_IMMUTABLE_SIMULATOR_SET: [u8; 4] = hex!("ad7e232e"); + /// Represents the context for [CheatcodeContext] #[derive(Debug, Default)] pub struct CheatcodeTracerContext<'a> { @@ -88,7 +95,10 @@ pub struct CheatcodeTracerContext<'a> { /// Tracer result to return back to foundry. #[derive(Debug, Default)] pub struct CheatcodeTracerResult { + /// Recorded expected calls. pub expected_calls: ExpectedCallTracker, + /// Immutables recorded via calls to ImmutableSimulator::setImmutables. + pub recorded_immutables: HashMap>>, } /// Defines the context for a Vm call. @@ -115,7 +125,7 @@ pub struct CallContext { pub is_static: bool, /// L1 block hashes to return when `BLOCKHASH` opcode is encountered. This ensures consistency /// when returning environment data in L2. - pub block_hashes: HashMap>, + pub block_hashes: HashMap>, } /// A tracer to allow for foundry-specific functionality. @@ -131,6 +141,8 @@ pub struct CheatcodeTracer { pub result: Arc>, /// Handle farcall state. farcall_handler: FarCallHandler, + /// Immutables recorded via calls to ImmutableSimulator::setImmutables. + recorded_immutables: HashMap>>, } impl CheatcodeTracer { @@ -360,6 +372,54 @@ impl DynTracer> for Cheatcode } } + // record immutables for an address during creates + if self.call_context.is_create { + if let Opcode::FarCall(_call) = data.opcode.variant.opcode { + let calldata = get_calldata(&state, memory); + let current = state.vm_local_state.callstack.current; + + if current.code_address == IMMUTABLE_SIMULATOR_STORAGE_ADDRESS && + calldata.starts_with(&SELECTOR_IMMUTABLE_SIMULATOR_SET) + { + let mut params = ethabi::decode( + &[ + ethabi::ParamType::Address, + ethabi::ParamType::Array(Box::new(ethabi::ParamType::Tuple(vec![ + ethabi::ParamType::Uint(256), + ethabi::ParamType::FixedBytes(32), + ]))), + ], + &calldata[4..], + ) + .expect("failed decoding setImmutables parameters"); + + let address = params.remove(0).into_address().expect("must be valid address"); + let immutables = params.remove(0).into_array().expect("must be valid array"); + for immutable in immutables { + let mut imm_tuple = immutable.into_tuple().expect("must be valid tuple"); + let imm_index = + imm_tuple.remove(0).into_uint().expect("must be valid uint").to_ru256(); + let imm_value = imm_tuple + .remove(0) + .into_fixed_bytes() + .expect("must be valid fixed bytes"); + let imm_value = FixedBytes::<32>::from_slice(&imm_value); + + self.recorded_immutables + .entry(address) + .and_modify(|entry| { + entry.insert(imm_index, imm_value); + }) + .or_insert_with(|| { + let mut value = HashMap::default(); + value.insert(imm_index, imm_value); + value + }); + } + } + } + } + if let Some(delegate_as) = self.call_context.delegate_as { if let Opcode::FarCall(_call) = data.opcode.variant.opcode { let current = state.vm_local_state.callstack.current; @@ -404,7 +464,11 @@ impl VmTracer for CheatcodeTracer { _stop_reason: zksync_multivm::interface::tracer::VmExecutionStopReason, ) { let cell = self.result.as_ref(); - cell.set(CheatcodeTracerResult { expected_calls: self.expected_calls.clone() }).unwrap(); + cell.set(CheatcodeTracerResult { + expected_calls: self.expected_calls.clone(), + recorded_immutables: self.recorded_immutables.clone(), + }) + .unwrap(); } } diff --git a/testdata/zk/ForkImmutableVars.sol b/testdata/zk/ForkImmutableVars.sol new file mode 100644 index 000000000..3e4bbd951 --- /dev/null +++ b/testdata/zk/ForkImmutableVars.sol @@ -0,0 +1,46 @@ +// SPDX-License-Identifier: UNLICENSED +pragma solidity ^0.8.13; + +import "ds-test/test.sol"; +import "../cheats/Vm.sol"; +import {Globals} from "./Globals.sol"; + +contract Counter { + uint256 public immutable SOME_IMMUTABLE_VARIABLE; + + constructor(uint256 value) { + SOME_IMMUTABLE_VARIABLE = value; + } + + uint256 public a; + + function set(uint256 _a) public { + a = _a; + } + + function get() public view returns (uint256) { + return a; + } +} + +contract ZkForkImmutableVarsTest is DSTest { + Vm constant vm = Vm(HEVM_ADDRESS); + uint256 constant ERA_FORK_BLOCK = 19579636; + uint256 constant IMMUTABLE_VAR_VALUE = 0xdeadbeef; + + function setUp() public { + vm.createSelectFork(Globals.ZKSYNC_MAINNET_URL, ERA_FORK_BLOCK); + } + + function testZkImmutableVariablesPersistedAfterFork() public { + Counter counter = new Counter(IMMUTABLE_VAR_VALUE); + assertEq(IMMUTABLE_VAR_VALUE, counter.SOME_IMMUTABLE_VARIABLE()); + + vm.makePersistent(address(counter)); + assertTrue(vm.isPersistent(address(counter))); + + vm.createSelectFork(Globals.ZKSYNC_MAINNET_URL, ERA_FORK_BLOCK - 100); + + assertEq(IMMUTABLE_VAR_VALUE, counter.SOME_IMMUTABLE_VARIABLE()); + } +}