diff --git a/crates/cheatcodes/src/strategy.rs b/crates/cheatcodes/src/strategy.rs index f1a37063b..63a01a2a4 100644 --- a/crates/cheatcodes/src/strategy.rs +++ b/crates/cheatcodes/src/strategy.rs @@ -19,6 +19,8 @@ use crate::{ }; pub trait CheatcodeInspectorStrategy: Debug + Send + Sync { + fn name(&self) -> &'static str; + /// Get nonce. fn get_nonce(&mut self, ccx: &mut CheatsCtxt, address: Address) -> Result { let account = ccx.ecx.journaled_state.load_account(address, &mut ccx.ecx.db)?; @@ -179,11 +181,11 @@ pub trait CheatcodeInspectorStrategy: Debug + Send + Sync { /// We define this in our fork pub trait CheatcodeInspectorStrategyExt: CheatcodeInspectorStrategy { - fn zksync_skip_zkvm(&mut self) -> Result { + fn zksync_cheatcode_skip_zkvm(&mut self) -> Result { unimplemented!() } - fn zksync_set_paymaster( + fn zksync_cheatcode_set_paymaster( &mut self, _paymaster_address: Address, _paymaster_input: &Bytes, @@ -191,11 +193,11 @@ pub trait CheatcodeInspectorStrategyExt: CheatcodeInspectorStrategy { unimplemented!() } - fn zksync_use_factory_deps(&mut self, _name: String) -> Result { + fn zksync_cheatcode_use_factory_deps(&mut self, _name: String) -> Result { unimplemented!() } - fn zksync_register_contract( + fn zksync_cheatcode_register_contract( &mut self, _name: String, _zk_bytecode_hash: FixedBytes<32>, @@ -208,17 +210,15 @@ pub trait CheatcodeInspectorStrategyExt: CheatcodeInspectorStrategy { unimplemented!() } - fn zksync_record_create_address(&mut self, _outcome: &CreateOutcome) { + fn zksync_cheatcode_select_zk_vm(&mut self, _data: InnerEcx, _enable: bool) { unimplemented!() } - fn zksync_sync_nonce(&mut self, _sender: Address, _nonce: u64, _ecx: Ecx) { - unimplemented!() - } + fn zksync_record_create_address(&mut self, _outcome: &CreateOutcome) {} - fn zksync_set_deployer_call_input(&mut self, _call: &mut CallInputs) { - unimplemented!() - } + fn zksync_sync_nonce(&mut self, _sender: Address, _nonce: u64, _ecx: Ecx) {} + + fn zksync_set_deployer_call_input(&mut self, _call: &mut CallInputs) {} fn zksync_try_create( &mut self, @@ -227,7 +227,7 @@ pub trait CheatcodeInspectorStrategyExt: CheatcodeInspectorStrategy { _input: &dyn CommonCreateInput, _executor: &mut dyn CheatcodesExecutor, ) -> Option { - unimplemented!() + None } fn zksync_try_call( @@ -237,22 +237,20 @@ pub trait CheatcodeInspectorStrategyExt: CheatcodeInspectorStrategy { _input: &CallInputs, _executor: &mut dyn CheatcodesExecutor, ) -> Option { - unimplemented!() + None } - fn zksync_select_fork_vm(&mut self, _data: InnerEcx, _fork_id: LocalForkId) { - unimplemented!() - } - - fn zksync_select_zk_vm(&mut self, _data: InnerEcx, _enable: bool) { - unimplemented!() - } + fn zksync_select_fork_vm(&mut self, _data: InnerEcx, _fork_id: LocalForkId) {} } #[derive(Debug, Default, Clone)] pub struct EvmCheatcodeInspectorStrategy; impl CheatcodeInspectorStrategy for EvmCheatcodeInspectorStrategy { + fn name(&self) -> &'static str { + "evm" + } + fn record_broadcastable_create_transactions( &mut self, _config: Arc, diff --git a/crates/cheatcodes/src/test.rs b/crates/cheatcodes/src/test.rs index 3cbaad954..c23a4f6ff 100644 --- a/crates/cheatcodes/src/test.rs +++ b/crates/cheatcodes/src/test.rs @@ -19,7 +19,7 @@ impl Cheatcode for zkVmCall { .strategy .lock() .expect("failed acquiring strategy") - .zksync_select_zk_vm(ccx.ecx, enable); + .zksync_cheatcode_select_zk_vm(ccx.ecx, enable); Ok(Default::default()) } @@ -27,7 +27,7 @@ impl Cheatcode for zkVmCall { impl Cheatcode for zkVmSkipCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { - ccx.state.strategy.lock().expect("failed acquiring strategy").zksync_skip_zkvm() + ccx.state.strategy.lock().expect("failed acquiring strategy").zksync_cheatcode_skip_zkvm() } } @@ -38,7 +38,7 @@ impl Cheatcode for zkUsePaymasterCall { .strategy .lock() .expect("failed acquiring strategy") - .zksync_set_paymaster(*paymaster_address, paymaster_input) + .zksync_cheatcode_set_paymaster(*paymaster_address, paymaster_input) } } @@ -49,7 +49,7 @@ impl Cheatcode for zkUseFactoryDepCall { .strategy .lock() .expect("failed acquiring strategy") - .zksync_use_factory_deps(name.clone()) + .zksync_cheatcode_use_factory_deps(name.clone()) } } @@ -64,15 +64,19 @@ impl Cheatcode for zkRegisterContractCall { zkDeployedBytecode, } = self; - ccx.state.strategy.lock().expect("failed acquiring strategy").zksync_register_contract( - name.clone(), - zkBytecodeHash.0.into(), - zkDeployedBytecode.to_vec(), - vec![], //TODO: add argument to cheatcode - *evmBytecodeHash, - evmDeployedBytecode.to_vec(), - evmBytecode.to_vec(), - ) + ccx.state + .strategy + .lock() + .expect("failed acquiring strategy") + .zksync_cheatcode_register_contract( + name.clone(), + zkBytecodeHash.0.into(), + zkDeployedBytecode.to_vec(), + vec![], //TODO: add argument to cheatcode + *evmBytecodeHash, + evmDeployedBytecode.to_vec(), + evmBytecode.to_vec(), + ) } } diff --git a/crates/chisel/src/executor.rs b/crates/chisel/src/executor.rs index 677f909b8..446d0b878 100644 --- a/crates/chisel/src/executor.rs +++ b/crates/chisel/src/executor.rs @@ -15,7 +15,7 @@ use foundry_evm::{ backend::Backend, decode::decode_console_logs, executors::{ - strategy::{EvmExecutorStrategy, ExecutorStrategy}, + strategy::{EvmExecutorStrategy, ExecutorStrategyExt}, ExecutorBuilder, }, inspectors::CheatsConfig, @@ -325,7 +325,7 @@ impl SessionSource { let env = self.config.evm_opts.evm_env().await.expect("Could not instantiate fork environment"); - let executor_strategy: Arc> = + let executor_strategy: Arc> = if self.config.foundry_config.zksync.run_in_zk_mode() { Arc::new(Mutex::new(ZksyncExecutorStrategy::default())) } else { @@ -362,7 +362,7 @@ impl SessionSource { executor_strategy .lock() .expect("failed acquiring strategy") - .new_cheatcode_inspector_strategy(Default::default()), + .new_cheatcode_inspector_strategy(), ) .into(), ) diff --git a/crates/cli/src/utils/mod.rs b/crates/cli/src/utils/mod.rs index 56fdec6d0..217b233bd 100644 --- a/crates/cli/src/utils/mod.rs +++ b/crates/cli/src/utils/mod.rs @@ -9,7 +9,7 @@ use foundry_common::{ shell, }; use foundry_config::{Chain, Config}; -use foundry_evm::executors::strategy::{EvmExecutorStrategy, ExecutorStrategy}; +use foundry_evm::executors::strategy::{EvmExecutorStrategy, ExecutorStrategyExt}; use foundry_strategy_zksync::ZksyncExecutorStrategy; use serde::de::DeserializeOwned; use std::{ @@ -94,7 +94,7 @@ pub fn get_provider(config: &Config) -> Result { get_provider_builder(config)?.build() } -pub fn get_executor_strategy(config: &Config) -> Arc> { +pub fn get_executor_strategy(config: &Config) -> Arc> { if config.zksync.run_in_zk_mode() { Arc::new(Mutex::new(ZksyncExecutorStrategy::default())) } else { diff --git a/crates/evm/core/src/backend/cow.rs b/crates/evm/core/src/backend/cow.rs index 06d1f305b..5b3cf2d42 100644 --- a/crates/evm/core/src/backend/cow.rs +++ b/crates/evm/core/src/backend/cow.rs @@ -1,6 +1,6 @@ //! A wrapper around `Backend` that is clone-on-write used for fuzzing. -use super::{BackendError, ForkInfo}; +use super::{strategy::BackendStrategyExt, BackendError, ForkInfo}; use crate::{ backend::{ diagnostic::RevertDiagnostic, Backend, DatabaseExt, LocalForkId, RevertStateSnapshotAction, @@ -22,7 +22,8 @@ use revm::{ }; use std::{ borrow::Cow, - collections::{BTreeMap, HashSet}, + collections::BTreeMap, + sync::{Arc, Mutex}, }; /// A wrapper around `Backend` that ensures only `revm::DatabaseRef` functions are called. @@ -118,8 +119,8 @@ 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 get_strategy(&mut self) -> Arc> { + self.backend.as_ref().strategy.clone() } fn snapshot_state(&mut self, journaled_state: &JournaledState, env: &Env) -> U256 { diff --git a/crates/evm/core/src/backend/mod.rs b/crates/evm/core/src/backend/mod.rs index 3f117e81b..b1013dc43 100644 --- a/crates/evm/core/src/backend/mod.rs +++ b/crates/evm/core/src/backend/mod.rs @@ -31,7 +31,7 @@ use std::{ sync::{Arc, Mutex}, time::Instant, }; -use strategy::{BackendStrategy, BackendStrategyForkInfo}; +use strategy::{BackendStrategyExt, BackendStrategyForkInfo}; mod diagnostic; pub use diagnostic::RevertDiagnostic; @@ -102,13 +102,8 @@ 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); + /// Retrieve the strategy. + fn get_strategy(&mut self) -> Arc>; /// Reverts the snapshot if it exists /// @@ -466,7 +461,7 @@ struct _ObjectSafe(dyn DatabaseExt); #[must_use] pub struct Backend { /// The behavior strategy. - pub strategy: Arc>, + pub strategy: Arc>, /// The access point for managing forks forks: MultiFork, @@ -497,15 +492,6 @@ pub struct Backend { inner: BackendInner, /// Keeps track of the fork type fork_url_type: CachedForkType, - /// TODO: Ensure this parameter is updated on `select_fork`. - /// - /// Keeps track if the backend is in ZK mode. - /// This is required to correctly merge storage when selecting another ZK fork. - /// 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 { @@ -513,7 +499,7 @@ impl Backend { /// /// If `fork` is `Some` this will use a `fork` database, otherwise with an in-memory /// database. - pub fn spawn(fork: Option, strategy: Arc>) -> Self { + pub fn spawn(fork: Option, strategy: Arc>) -> Self { Self::new(MultiFork::spawn(), fork, strategy) } @@ -526,7 +512,7 @@ impl Backend { pub fn new( forks: MultiFork, fork: Option, - strategy: Arc>, + strategy: Arc>, ) -> Self { trace!(target: "backend", forking_mode=?fork.is_some(), "creating executor backend"); // Note: this will take of registering the `fork` @@ -542,8 +528,6 @@ impl Backend { active_fork_ids: None, inner, fork_url_type: Default::default(), - is_zk: false, - zk_recorded_immutable_keys: Default::default(), strategy, }; @@ -571,7 +555,7 @@ impl Backend { id: &ForkId, fork: Fork, journaled_state: JournaledState, - strategy: Arc>, + strategy: Arc>, ) -> Self { let mut backend = Self::spawn(None, strategy); let fork_ids = backend.inner.insert_new_fork(id.clone(), fork.db, journaled_state); @@ -589,8 +573,6 @@ impl Backend { active_fork_ids: None, inner: Default::default(), fork_url_type: Default::default(), - is_zk: false, - zk_recorded_immutable_keys: Default::default(), strategy: self.strategy.clone(), } } @@ -996,11 +978,8 @@ 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 get_strategy(&mut self) -> Arc> { + self.strategy.clone() } fn snapshot_state(&mut self, journaled_state: &JournaledState, env: &Env) -> U256 { @@ -1875,7 +1854,7 @@ impl BackendInner { id: LocalForkId, new_fork_id: ForkId, backend: SharedBackend, - strategy: Arc>, + strategy: Arc>, ) -> eyre::Result { let fork_id = self.ensure_fork_id(id)?; let idx = self.ensure_fork_index(fork_id)?; @@ -2007,7 +1986,7 @@ fn commit_transaction( fork_id: &ForkId, persistent_accounts: &HashSet
, inspector: &mut dyn InspectorExt, - strategy: Arc>, + strategy: Arc>, ) -> eyre::Result<()> { // TODO: Remove after https://github.com/foundry-rs/foundry/pull/9131 // if the tx has the blob_versioned_hashes field, we assume it's a Cancun block diff --git a/crates/evm/core/src/backend/strategy.rs b/crates/evm/core/src/backend/strategy.rs index c1dd18ccc..c9adab95f 100644 --- a/crates/evm/core/src/backend/strategy.rs +++ b/crates/evm/core/src/backend/strategy.rs @@ -3,12 +3,12 @@ use std::fmt::Debug; use crate::InspectorExt; use super::{BackendInner, DatabaseExt, Fork, ForkDB, ForkType, FoundryEvmInMemoryDB}; -use alloy_primitives::Address; +use alloy_primitives::{Address, U256}; use alloy_rpc_types::serde_helpers::OtherFields; use eyre::Context; use revm::{ db::CacheDB, - primitives::{EnvWithHandlerCfg, ResultAndState}, + primitives::{EnvWithHandlerCfg, HashSet, ResultAndState}, DatabaseRef, JournaledState, }; use serde::{Deserialize, Serialize}; @@ -20,10 +20,6 @@ pub struct BackendStrategyForkInfo<'a> { } pub trait BackendStrategy: Debug + Send + Sync { - // fn new() -> Arc> { - // Arc::new(Mutex::new(Self::default())) - // } - fn name(&self) -> &'static str; /// When creating or switching forks, we update the AccountInfo of the contract @@ -90,6 +86,16 @@ pub trait BackendStrategy: Debug + Send + Sync { } } +pub trait BackendStrategyExt: BackendStrategy { + /// 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 zksync_save_immutable_storage(&mut self, _addr: Address, _keys: HashSet) {} +} + struct _ObjectSafe(dyn BackendStrategy); #[derive(Debug, Default, Clone, Serialize, Deserialize)] @@ -137,6 +143,8 @@ impl BackendStrategy for EvmBackendStrategy { fn set_inspect_context(&mut self, _other_fields: OtherFields) {} } +impl BackendStrategyExt for EvmBackendStrategy {} + impl EvmBackendStrategy { /// Merges the state of all `accounts` from the currently active db into the given `fork` pub(crate) fn update_fork_db_contracts( diff --git a/crates/evm/evm/src/executors/builder.rs b/crates/evm/evm/src/executors/builder.rs index 5dd29ace9..22eb0e377 100644 --- a/crates/evm/evm/src/executors/builder.rs +++ b/crates/evm/evm/src/executors/builder.rs @@ -4,7 +4,7 @@ use crate::{executors::Executor, inspectors::InspectorStackBuilder}; use foundry_evm_core::backend::Backend; use revm::primitives::{Env, EnvWithHandlerCfg, SpecId}; -use super::strategy::ExecutorStrategy; +use super::strategy::ExecutorStrategyExt; /// The builder that allows to configure an evm [`Executor`] which a stack of optional /// [`revm::Inspector`]s, such as [`Cheatcodes`]. @@ -82,7 +82,7 @@ impl ExecutorBuilder { self, env: Env, db: Backend, - strategy: Arc>, + strategy: Arc>, ) -> Executor { let Self { mut stack, gas_limit, spec_id, legacy_assertions } = self; if stack.block.is_none() { diff --git a/crates/evm/evm/src/executors/mod.rs b/crates/evm/evm/src/executors/mod.rs index 25af18dc1..67e653a72 100644 --- a/crates/evm/evm/src/executors/mod.rs +++ b/crates/evm/evm/src/executors/mod.rs @@ -40,7 +40,7 @@ use std::{ borrow::Cow, sync::{Arc, Mutex}, }; -use strategy::ExecutorStrategy; +use strategy::ExecutorStrategyExt; mod builder; pub use builder::ExecutorBuilder; @@ -96,7 +96,7 @@ pub struct Executor { /// Whether `failed()` should be called on the test contract to determine if the test failed. legacy_assertions: bool, - strategy: Arc>, + strategy: Arc>, } impl Executor { @@ -114,7 +114,7 @@ impl Executor { inspector: InspectorStack, gas_limit: u64, legacy_assertions: bool, - strategy: Arc>, + strategy: Arc>, ) -> Self { // Need to create a non-empty contract on the cheatcodes address so `extcodesize` checks // do not fail. diff --git a/crates/evm/evm/src/executors/strategy.rs b/crates/evm/evm/src/executors/strategy.rs index 0be7c72fb..12e2ef85b 100644 --- a/crates/evm/evm/src/executors/strategy.rs +++ b/crates/evm/evm/src/executors/strategy.rs @@ -6,7 +6,7 @@ use std::{ use alloy_primitives::{Address, U256}; use foundry_cheatcodes::strategy::{CheatcodeInspectorStrategyExt, EvmCheatcodeInspectorStrategy}; use foundry_evm_core::backend::{ - strategy::{BackendStrategy, EvmBackendStrategy}, + strategy::{BackendStrategyExt, EvmBackendStrategy}, BackendResult, }; use foundry_zksync_compiler::DualCompiledContracts; @@ -15,6 +15,8 @@ use revm::DatabaseRef; use super::Executor; pub trait ExecutorStrategy: Debug + Send + Sync { + fn name(&self) -> &'static str; + fn set_balance( &mut self, executor: &mut Executor, @@ -29,19 +31,28 @@ pub trait ExecutorStrategy: Debug + Send + Sync { nonce: u64, ) -> BackendResult<()>; - fn new_backend_strategy(&self) -> Arc>; - fn new_cheatcode_inspector_strategy( - &self, - dual_compiled_contracts: DualCompiledContracts, - ) -> Arc>; + fn new_backend_strategy(&self) -> Arc>; + fn new_cheatcode_inspector_strategy(&self) -> Arc>; // TODO perhaps need to create fresh strategies as well } +pub trait ExecutorStrategyExt: ExecutorStrategy { + fn zksync_set_dual_compiled_contracts( + &mut self, + _dual_compiled_contracts: DualCompiledContracts, + ) { + } +} + #[derive(Debug, Default, Clone)] pub struct EvmExecutorStrategy {} impl ExecutorStrategy for EvmExecutorStrategy { + fn name(&self) -> &'static str { + "evm" + } + fn set_balance( &mut self, executor: &mut Executor, @@ -69,14 +80,13 @@ impl ExecutorStrategy for EvmExecutorStrategy { Ok(()) } - fn new_backend_strategy(&self) -> Arc> { + fn new_backend_strategy(&self) -> Arc> { Arc::new(Mutex::new(EvmBackendStrategy)) } - fn new_cheatcode_inspector_strategy( - &self, - _dual_compiled_contracts: DualCompiledContracts, - ) -> Arc> { + fn new_cheatcode_inspector_strategy(&self) -> Arc> { Arc::new(Mutex::new(EvmCheatcodeInspectorStrategy)) } } + +impl ExecutorStrategyExt for EvmExecutorStrategy {} diff --git a/crates/evm/evm/src/executors/trace.rs b/crates/evm/evm/src/executors/trace.rs index be873c6ad..83a06a516 100644 --- a/crates/evm/evm/src/executors/trace.rs +++ b/crates/evm/evm/src/executors/trace.rs @@ -9,7 +9,7 @@ use std::{ sync::{Arc, Mutex}, }; -use super::strategy::ExecutorStrategy; +use super::strategy::ExecutorStrategyExt; /// A default executor with tracing enabled pub struct TracingExecutor { @@ -24,7 +24,7 @@ impl TracingExecutor { debug: bool, decode_internal: bool, alphanet: bool, - strategy: Arc>, + strategy: Arc>, ) -> Self { let db = Backend::spawn( fork, diff --git a/crates/forge/bin/cmd/coverage.rs b/crates/forge/bin/cmd/coverage.rs index c4b6451dc..e7ea405a3 100644 --- a/crates/forge/bin/cmd/coverage.rs +++ b/crates/forge/bin/cmd/coverage.rs @@ -13,14 +13,13 @@ use forge::{ utils::IcPcMap, MultiContractRunnerBuilder, TestOptions, }; -use foundry_cli::utils::{LoadConfig, STATIC_FUZZ_SEED}; +use foundry_cli::utils::{self, LoadConfig, STATIC_FUZZ_SEED}; use foundry_common::{compile::ProjectCompiler, fs}; use foundry_compilers::{ artifacts::{sourcemap::SourceMap, CompactBytecode, CompactDeployedBytecode, SolcLanguage}, Artifact, ArtifactId, Project, ProjectCompileOutput, }; use foundry_config::{Config, SolcReq}; -use foundry_zksync_compiler::DualCompiledContracts; use rayon::prelude::*; use semver::Version; use std::{ @@ -223,6 +222,7 @@ impl CoverageArgs { ) -> Result<()> { let root = project.paths.root; let verbosity = evm_opts.verbosity; + let strategy = utils::get_executor_strategy(&config); // Build the contract runner let env = evm_opts.evm_env().await?; @@ -237,7 +237,7 @@ impl CoverageArgs { ..Default::default() }) .set_coverage(true) - .build(&root, output.clone(), None, env, evm_opts, DualCompiledContracts::default())?; + .build(&root, output.clone(), None, env, evm_opts, strategy)?; let known_contracts = runner.known_contracts.clone(); diff --git a/crates/forge/bin/cmd/test/mod.rs b/crates/forge/bin/cmd/test/mod.rs index d1d510bc4..6139dc71d 100644 --- a/crates/forge/bin/cmd/test/mod.rs +++ b/crates/forge/bin/cmd/test/mod.rs @@ -271,6 +271,7 @@ impl TestArgs { pub async fn execute_tests(mut self) -> Result { // Merge all configs. let (mut config, mut evm_opts) = self.load_config_and_evm_opts_emit_warnings()?; + let strategy = utils::get_executor_strategy(&config); // Explicitly enable isolation for gas reports for more correct gas accounting. if self.gas_report { @@ -367,6 +368,11 @@ impl TestArgs { // Prepare the test builder. let config = Arc::new(config); + strategy + .lock() + .expect("failed acquiring strategy") + .zksync_set_dual_compiled_contracts(dual_compiled_contracts.unwrap_or_default()); + let runner = MultiContractRunnerBuilder::new(config.clone()) .set_debug(should_debug) .set_decode_internal(decode_internal) @@ -377,14 +383,7 @@ impl TestArgs { .with_test_options(test_options.clone()) .enable_isolation(evm_opts.isolate) .alphanet(evm_opts.alphanet) - .build( - project_root, - output.clone(), - zk_output, - env, - evm_opts, - dual_compiled_contracts.unwrap_or_default(), - )?; + .build(project_root, output.clone(), zk_output, env, evm_opts, strategy)?; let mut maybe_override_mt = |flag, maybe_regex: Option<&Option>| { if let Some(Some(regex)) = maybe_regex { diff --git a/crates/forge/src/multi_runner.rs b/crates/forge/src/multi_runner.rs index f4f7b907d..2ea177733 100644 --- a/crates/forge/src/multi_runner.rs +++ b/crates/forge/src/multi_runner.rs @@ -7,7 +7,6 @@ use crate::{ use alloy_json_abi::{Function, JsonAbi}; use alloy_primitives::{Address, Bytes, U256}; use eyre::Result; -use foundry_cli::utils; use foundry_common::{get_contract_name, ContractsByArtifact, TestFunctionExt}; use foundry_compilers::{ artifacts::{CompactBytecode, CompactContractBytecode, CompactDeployedBytecode, Libraries}, @@ -19,7 +18,7 @@ use foundry_config::Config; use foundry_evm::{ backend::Backend, decode::RevertDecoder, - executors::{strategy::ExecutorStrategy, ExecutorBuilder}, + executors::{strategy::ExecutorStrategyExt, ExecutorBuilder}, fork::CreateFork, inspectors::CheatsConfig, opts::EvmOpts, @@ -27,7 +26,6 @@ use foundry_evm::{ traces::{InternalTraceMode, TraceMode}, }; use foundry_linking::{LinkOutput, Linker}; -use foundry_zksync_compiler::DualCompiledContracts; use rayon::prelude::*; use revm::primitives::SpecId; @@ -86,10 +84,8 @@ pub struct MultiContractRunner { pub libs_to_deploy: Vec, /// Library addresses used to link contracts. pub libraries: Libraries, - /// Dual compiled contracts - pub dual_compiled_contracts: DualCompiledContracts, - /// Use zk runner. - pub use_zk: bool, + /// Execution strategy. + pub strategy: Arc>, } impl MultiContractRunner { @@ -182,12 +178,10 @@ impl MultiContractRunner { trace!("running all tests"); // The DB backend that serves all the data. - let strategy = utils::get_executor_strategy(&self.config); - let mut db = Backend::spawn( + let db = Backend::spawn( self.fork.take(), - strategy.lock().expect("failed acquiring strategy").new_backend_strategy(), + self.strategy.lock().expect("failed acquiring strategy").new_backend_strategy(), ); - db.is_zk = self.use_zk; let find_timer = Instant::now(); let contracts = self.matching_contracts(filter).collect::>(); @@ -215,7 +209,6 @@ impl MultiContractRunner { filter, &tokio_handle, Some(&tests_progress), - strategy.clone(), ); tests_progress @@ -235,15 +228,8 @@ impl MultiContractRunner { } else { contracts.par_iter().for_each(|&(id, contract)| { let _guard = tokio_handle.enter(); - let result = self.run_test_suite( - id, - contract, - db.clone(), - filter, - &tokio_handle, - None, - strategy.clone(), - ); + let result = + self.run_test_suite(id, contract, db.clone(), filter, &tokio_handle, None); let _ = tx.send((id.identifier(), result)); }) } @@ -257,7 +243,6 @@ impl MultiContractRunner { filter: &dyn TestFilter, tokio_handle: &tokio::runtime::Handle, progress: Option<&TestsProgress>, - strategy: Arc>, ) -> SuiteResult { let identifier = artifact_id.identifier(); let mut span_name = identifier.as_str(); @@ -268,10 +253,10 @@ impl MultiContractRunner { Some(self.known_contracts.clone()), Some(artifact_id.name.clone()), Some(artifact_id.version.clone()), - strategy + self.strategy .lock() .expect("failed acquiring strategy") - .new_cheatcode_inspector_strategy(self.dual_compiled_contracts.clone()), + .new_cheatcode_inspector_strategy(), ); let trace_mode = TraceMode::default() @@ -291,7 +276,7 @@ impl MultiContractRunner { .spec(self.evm_spec) .gas_limit(self.evm_opts.gas_limit()) .legacy_assertions(self.config.legacy_assertions) - .build(self.env.clone(), db, strategy); + .build(self.env.clone(), db, self.strategy.clone()); if !enabled!(tracing::Level::TRACE) { span_name = get_contract_name(&identifier); @@ -428,9 +413,8 @@ impl MultiContractRunnerBuilder { zk_output: Option, env: revm::primitives::Env, evm_opts: EvmOpts, - dual_compiled_contracts: DualCompiledContracts, + strategy: Arc>, ) -> Result { - let use_zk = zk_output.is_some(); let mut known_contracts = ContractsByArtifact::default(); let output = output.with_stripped_file_prefixes(root); let linker = Linker::new(root, output.artifact_ids().collect()); @@ -472,7 +456,7 @@ impl MultiContractRunnerBuilder { } } - if !use_zk { + if zk_output.is_none() { known_contracts = ContractsByArtifact::new(linked_contracts); } else if let Some(zk_output) = zk_output { let zk_contracts = zk_output.with_stripped_file_prefixes(root).into_artifacts(); @@ -530,8 +514,7 @@ impl MultiContractRunnerBuilder { known_contracts, libs_to_deploy, libraries, - dual_compiled_contracts, - use_zk, + strategy, }) } } diff --git a/crates/forge/tests/it/test_helpers.rs b/crates/forge/tests/it/test_helpers.rs index fe06d785f..6a1b0eef6 100644 --- a/crates/forge/tests/it/test_helpers.rs +++ b/crates/forge/tests/it/test_helpers.rs @@ -3,9 +3,10 @@ use alloy_chains::NamedChain; use alloy_primitives::U256; use forge::{ - revm::primitives::SpecId, MultiContractRunner, MultiContractRunnerBuilder, TestOptions, - TestOptionsBuilder, + executors::strategy::EvmExecutorStrategy, revm::primitives::SpecId, MultiContractRunner, + MultiContractRunnerBuilder, TestOptions, TestOptionsBuilder, }; +use foundry_cli::utils; use foundry_compilers::{ artifacts::{EvmVersion, Libraries, Settings}, utils::RuntimeOrHandle, @@ -35,7 +36,7 @@ use std::{ env, fmt, io::Write, path::{Path, PathBuf}, - sync::{Arc, LazyLock}, + sync::{Arc, LazyLock, Mutex}, }; type ZkProject = Project; @@ -322,13 +323,14 @@ impl ForgeTestData { let sender = config.sender; + let strategy = utils::get_executor_strategy(&config); let mut builder = self.base_runner(); builder.config = Arc::new(config); builder .enable_isolation(opts.isolate) .sender(sender) .with_test_options(self.test_opts.clone()) - .build(root, output, None, env, opts, Default::default()) + .build(root, output, None, env, opts, strategy) .unwrap() } @@ -356,13 +358,18 @@ impl ForgeTestData { test_opts.fuzz.no_zksync_reserved_addresses = zk_config.fuzz.no_zksync_reserved_addresses; let sender = zk_config.sender; + let strategy = utils::get_executor_strategy(&zk_config); + strategy + .lock() + .expect("failed acquiring strategy") + .zksync_set_dual_compiled_contracts(dual_compiled_contracts); let mut builder = self.base_runner(); builder.config = Arc::new(zk_config); builder .enable_isolation(opts.isolate) .sender(sender) .with_test_options(test_opts) - .build(root, output, Some(zk_output), env, opts, dual_compiled_contracts) + .build(root, output, Some(zk_output), env, opts, strategy) .unwrap() } @@ -377,7 +384,7 @@ impl ForgeTestData { None, opts.local_evm_env(), opts, - Default::default(), + Arc::new(Mutex::new(EvmExecutorStrategy::default())), ) .unwrap() } @@ -394,7 +401,14 @@ impl ForgeTestData { self.base_runner() .with_fork(fork) - .build(self.project.root(), self.output.clone(), None, env, opts, Default::default()) + .build( + self.project.root(), + self.output.clone(), + None, + env, + opts, + Arc::new(Mutex::new(EvmExecutorStrategy::default())), + ) .unwrap() } } diff --git a/crates/script/src/lib.rs b/crates/script/src/lib.rs index da51831b5..c5d602f38 100644 --- a/crates/script/src/lib.rs +++ b/crates/script/src/lib.rs @@ -630,6 +630,11 @@ impl ScriptConfig { if let Some((known_contracts, script_wallets, target, dual_compiled_contracts)) = cheats_data { + strategy + .lock() + .expect("failed acquiring strategy") + .zksync_set_dual_compiled_contracts(dual_compiled_contracts); + builder = builder.inspectors(|stack| { stack .cheatcodes( @@ -642,7 +647,7 @@ impl ScriptConfig { strategy .lock() .expect("failed acquiring strategy") - .new_cheatcode_inspector_strategy(dual_compiled_contracts), + .new_cheatcode_inspector_strategy(), ) .into(), ) diff --git a/crates/strategy/zksync/src/backend.rs b/crates/strategy/zksync/src/backend.rs index 1cc2b1b71..829707607 100644 --- a/crates/strategy/zksync/src/backend.rs +++ b/crates/strategy/zksync/src/backend.rs @@ -2,6 +2,7 @@ use std::collections::hash_map::Entry; use alloy_primitives::{map::HashMap, Address, U256}; use alloy_rpc_types::serde_helpers::OtherFields; +use foundry_evm::backend::strategy::BackendStrategyExt; use foundry_evm_core::{ backend::{ strategy::{ @@ -29,6 +30,7 @@ pub struct ZksyncBackendStrategy { evm: EvmBackendStrategy, inspect_context: Option, persisted_factory_deps: HashMap>, + /// Store storage keys per contract address for immutable variables. persistent_immutable_keys: HashMap>, } @@ -181,6 +183,15 @@ impl ZksyncBackendStrategy { } } +impl BackendStrategyExt for ZksyncBackendStrategy { + fn zksync_save_immutable_storage(&mut self, addr: Address, keys: HashSet) { + self.persistent_immutable_keys + .entry(addr) + .and_modify(|entry| entry.extend(&keys)) + .or_insert(keys); + } +} + pub(crate) struct ZksyncBackendMerge; /// Defines the zksync specific state to help during merge. diff --git a/crates/strategy/zksync/src/cheatcode.rs b/crates/strategy/zksync/src/cheatcode.rs index 75e2f11dc..d6c0f4ff5 100644 --- a/crates/strategy/zksync/src/cheatcode.rs +++ b/crates/strategy/zksync/src/cheatcode.rs @@ -211,6 +211,10 @@ impl ZkPersistNonceUpdate { } impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { + fn name(&self) -> &'static str { + "zk" + } + fn get_nonce(&mut self, ccx: &mut CheatsCtxt<'_, '_, '_, '_>, address: Address) -> Result { let nonce = foundry_zksync_core::nonce(address, ccx.ecx) as u64; Ok(nonce) @@ -615,12 +619,12 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { } impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { - fn zksync_skip_zkvm(&mut self) -> Result { + fn zksync_cheatcode_skip_zkvm(&mut self) -> Result { self.skip_zk_vm = true; Ok(Default::default()) } - fn zksync_set_paymaster( + fn zksync_cheatcode_set_paymaster( &mut self, paymaster_address: Address, paymaster_input: &Bytes, @@ -630,13 +634,13 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { Ok(Default::default()) } - fn zksync_use_factory_deps(&mut self, name: String) -> foundry_cheatcodes::Result { + fn zksync_cheatcode_use_factory_deps(&mut self, name: String) -> foundry_cheatcodes::Result { info!("Adding factory dependency: {:?}", name); self.zk_use_factory_deps.push(name); Ok(Default::default()) } - fn zksync_register_contract( + fn zksync_cheatcode_register_contract( &mut self, name: String, zk_bytecode_hash: FixedBytes<32>, @@ -866,8 +870,12 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { foundry_zksync_core::get_immutable_slot_key(addr, slot_index) .to_ru256() }) - .collect::>(); - ecx.db.save_zk_immutable_storage(addr, keys); + .collect::>(); + ecx.db + .get_strategy() + .lock() + .expect("failed acquiring strategy") + .zksync_save_immutable_storage(addr, keys); } } @@ -1092,7 +1100,7 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { self.select_fork_vm(data, fork_id); } - fn zksync_select_zk_vm(&mut self, data: InnerEcx<'_, '_, '_>, enable: bool) { + fn zksync_cheatcode_select_zk_vm(&mut self, data: InnerEcx<'_, '_, '_>, enable: bool) { if enable { self.select_zk_vm(data, None) } else { diff --git a/crates/strategy/zksync/src/executor.rs b/crates/strategy/zksync/src/executor.rs index fa46e01ac..0eeaa8d97 100644 --- a/crates/strategy/zksync/src/executor.rs +++ b/crates/strategy/zksync/src/executor.rs @@ -2,10 +2,11 @@ use std::sync::{Arc, Mutex}; use alloy_primitives::{Address, U256}; use foundry_cheatcodes::strategy::CheatcodeInspectorStrategyExt; + use foundry_evm::{ - backend::{strategy::BackendStrategy, BackendResult}, + backend::{strategy::BackendStrategyExt, BackendResult}, executors::{ - strategy::{EvmExecutorStrategy, ExecutorStrategy}, + strategy::{EvmExecutorStrategy, ExecutorStrategy, ExecutorStrategyExt}, Executor, }, }; @@ -17,9 +18,14 @@ use crate::{ZksyncBackendStrategy, ZksyncCheatcodeInspectorStrategy}; #[derive(Debug, Default, Clone)] pub struct ZksyncExecutorStrategy { evm: EvmExecutorStrategy, + dual_compiled_contracts: DualCompiledContracts, } impl ExecutorStrategy for ZksyncExecutorStrategy { + fn name(&self) -> &'static str { + "zk" + } + fn set_balance( &mut self, executor: &mut Executor, @@ -53,14 +59,22 @@ impl ExecutorStrategy for ZksyncExecutorStrategy { Ok(()) } - fn new_backend_strategy(&self) -> Arc> { + fn new_backend_strategy(&self) -> Arc> { Arc::new(Mutex::new(ZksyncBackendStrategy::default())) } - fn new_cheatcode_inspector_strategy( - &self, + fn new_cheatcode_inspector_strategy(&self) -> Arc> { + Arc::new(Mutex::new(ZksyncCheatcodeInspectorStrategy::new( + self.dual_compiled_contracts.clone(), + ))) + } +} + +impl ExecutorStrategyExt for ZksyncExecutorStrategy { + fn zksync_set_dual_compiled_contracts( + &mut self, dual_compiled_contracts: DualCompiledContracts, - ) -> Arc> { - Arc::new(Mutex::new(ZksyncCheatcodeInspectorStrategy::new(dual_compiled_contracts))) + ) { + self.dual_compiled_contracts = dual_compiled_contracts; } } diff --git a/crates/verify/src/utils.rs b/crates/verify/src/utils.rs index 54c402d85..51de9de19 100644 --- a/crates/verify/src/utils.rs +++ b/crates/verify/src/utils.rs @@ -16,7 +16,7 @@ use foundry_compilers::artifacts::{BytecodeHash, CompactContractBytecode, EvmVer use foundry_config::Config; use foundry_evm::{ constants::DEFAULT_CREATE2_DEPLOYER, - executors::{strategy::ExecutorStrategy, TracingExecutor}, + executors::{strategy::ExecutorStrategyExt, TracingExecutor}, opts::EvmOpts, }; use reqwest::Url; @@ -327,7 +327,7 @@ pub async fn get_tracing_executor( fork_blk_num: u64, evm_version: EvmVersion, evm_opts: EvmOpts, - strategy: Arc>, + strategy: Arc>, ) -> Result<(Env, TracingExecutor)> { fork_config.fork_block_number = Some(fork_blk_num); fork_config.evm_version = evm_version;