From 5ddccc2945108aaa531489c819adaceaadb3a99a Mon Sep 17 00:00:00 2001 From: Nisheeth Barthwal Date: Sat, 21 Dec 2024 22:02:46 +0100 Subject: [PATCH 1/6] fix re-entrancy --- crates/cheatcodes/src/config.rs | 10 +- crates/cheatcodes/src/evm.rs | 18 +- crates/cheatcodes/src/evm/fork.rs | 12 +- crates/cheatcodes/src/evm/mock.rs | 8 +- crates/cheatcodes/src/fs.rs | 2 +- crates/cheatcodes/src/inspector.rs | 86 ++--- crates/cheatcodes/src/lib.rs | 12 - crates/cheatcodes/src/strategy.rs | 154 +++++++-- crates/cheatcodes/src/test.rs | 42 ++- crates/chisel/src/executor.rs | 6 +- crates/cli/src/utils/mod.rs | 10 +- crates/evm/core/src/backend/cow.rs | 6 +- crates/evm/core/src/backend/mod.rs | 62 ++-- crates/evm/core/src/backend/strategy.rs | 72 +++- crates/evm/evm/src/executors/builder.rs | 4 +- crates/evm/evm/src/executors/mod.rs | 83 ++--- crates/evm/evm/src/executors/strategy.rs | 113 +++++-- crates/evm/evm/src/executors/trace.rs | 6 +- crates/forge/bin/cmd/test/mod.rs | 5 +- crates/forge/src/multi_runner.rs | 12 +- crates/forge/src/runner.rs | 6 +- crates/forge/tests/it/test_helpers.rs | 17 +- crates/script/src/lib.rs | 15 +- crates/script/src/runner.rs | 6 +- crates/strategy/zksync/src/backend.rs | 70 +++- crates/strategy/zksync/src/cheatcode.rs | 405 +++++++++++++++-------- crates/strategy/zksync/src/executor.rs | 127 +++++-- crates/strategy/zksync/src/lib.rs | 2 +- crates/verify/src/bytecode.rs | 4 +- crates/verify/src/utils.rs | 4 +- crates/zksync/compiler/src/zksolc/mod.rs | 5 + 31 files changed, 910 insertions(+), 474 deletions(-) diff --git a/crates/cheatcodes/src/config.rs b/crates/cheatcodes/src/config.rs index 00a393dde..9c519ee93 100644 --- a/crates/cheatcodes/src/config.rs +++ b/crates/cheatcodes/src/config.rs @@ -1,6 +1,6 @@ use super::Result; use crate::{ - strategy::{CheatcodeInspectorStrategy, EvmCheatcodeInspectorStrategy}, + strategy::{new_evm_strategy, Strategy}, Vm::Rpc, }; use alloy_primitives::{map::AddressHashMap, U256}; @@ -57,7 +57,7 @@ pub struct CheatsConfig { /// Version of the script/test contract which is currently running. pub running_version: Option, /// The behavior strategy. - pub strategy: Box, + pub strategy: Strategy, /// Whether to enable legacy (non-reverting) assertions. pub assertions_revert: bool, /// Optional seed for the RNG algorithm. @@ -73,7 +73,7 @@ impl CheatsConfig { available_artifacts: Option, running_contract: Option, running_version: Option, - strategy: Box, + strategy: Strategy, ) -> Self { let mut allowed_paths = vec![config.root.0.clone()]; allowed_paths.extend(config.libs.clone()); @@ -234,7 +234,7 @@ impl Default for CheatsConfig { available_artifacts: Default::default(), running_contract: Default::default(), running_version: Default::default(), - strategy: Box::new(EvmCheatcodeInspectorStrategy::default()), + strategy: new_evm_strategy(), assertions_revert: true, seed: None, } @@ -253,7 +253,7 @@ mod tests { None, None, None, - Box::new(EvmCheatcodeInspectorStrategy::default()), + new_evm_strategy(), ) } diff --git a/crates/cheatcodes/src/evm.rs b/crates/cheatcodes/src/evm.rs index 0b3b574ca..43eb9f50b 100644 --- a/crates/cheatcodes/src/evm.rs +++ b/crates/cheatcodes/src/evm.rs @@ -64,7 +64,7 @@ impl Cheatcode for getNonce_0Call { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { account } = self; - ccx.with_strategy(|strategy, ccx| strategy.cheatcode_get_nonce(ccx, *account)) + ccx.state.strategy.inner.clone().cheatcode_get_nonce(ccx, *account) } } @@ -350,7 +350,7 @@ impl Cheatcode for getBlobhashesCall { impl Cheatcode for rollCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { newHeight } = self; - ccx.with_strategy(|strategy, ccx| strategy.cheatcode_roll(ccx, *newHeight)) + ccx.state.strategy.inner.clone().cheatcode_roll(ccx, *newHeight) } } @@ -372,7 +372,7 @@ impl Cheatcode for txGasPriceCall { impl Cheatcode for warpCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { newTimestamp } = self; - ccx.with_strategy(|strategy, ccx| strategy.cheatcode_warp(ccx, *newTimestamp)) + ccx.state.strategy.inner.clone().cheatcode_warp(ccx, *newTimestamp) } } @@ -407,7 +407,7 @@ impl Cheatcode for dealCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { account: address, newBalance: new_balance } = *self; - ccx.with_strategy(|strategy, ccx| strategy.cheatcode_deal(ccx, address, new_balance)) + ccx.state.strategy.inner.clone().cheatcode_deal(ccx, address, new_balance) } } @@ -415,14 +415,14 @@ impl Cheatcode for etchCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { target, newRuntimeBytecode } = self; - ccx.with_strategy(|strategy, ccx| strategy.cheatcode_etch(ccx, *target, newRuntimeBytecode)) + ccx.state.strategy.inner.clone().cheatcode_etch(ccx, *target, newRuntimeBytecode) } } impl Cheatcode for resetNonceCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { account } = self; - ccx.with_strategy(|strategy, ccx| strategy.cheatcode_reset_nonce(ccx, *account)) + ccx.state.strategy.inner.clone().cheatcode_reset_nonce(ccx, *account) } } @@ -430,7 +430,7 @@ impl Cheatcode for setNonceCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { account, newNonce } = *self; - ccx.with_strategy(|strategy, ccx| strategy.cheatcode_set_nonce(ccx, account, newNonce)) + ccx.state.strategy.inner.clone().cheatcode_set_nonce(ccx, account, newNonce) } } @@ -438,9 +438,7 @@ impl Cheatcode for setNonceUnsafeCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { account, newNonce } = *self; - ccx.with_strategy(|strategy, ccx| { - strategy.cheatcode_set_nonce_unsafe(ccx, account, newNonce) - }) + ccx.state.strategy.inner.clone().cheatcode_set_nonce_unsafe(ccx, account, newNonce) } } diff --git a/crates/cheatcodes/src/evm/fork.rs b/crates/cheatcodes/src/evm/fork.rs index 17bb9bacf..a1bcc9ccb 100644 --- a/crates/cheatcodes/src/evm/fork.rs +++ b/crates/cheatcodes/src/evm/fork.rs @@ -125,7 +125,11 @@ impl Cheatcode for selectForkCall { persist_caller(ccx); check_broadcast(ccx.state)?; - ccx.with_strategy(|strategy, ccx| strategy.zksync_select_fork_vm(ccx.ecx, *forkId)); + ccx.state.strategy.inner.zksync_select_fork_vm( + ccx.state.strategy.context.as_mut(), + ccx.ecx, + *forkId, + ); ccx.ecx.db.select_fork(*forkId, &mut ccx.ecx.env, &mut ccx.ecx.journaled_state)?; Ok(Default::default()) @@ -281,7 +285,11 @@ fn create_select_fork(ccx: &mut CheatsCtxt, url_or_alias: &str, block: Option Result { let Self { callee, data, returnData } = self; - ccx.with_strategy(|strategy, ccx| { - strategy.cheatcode_mock_call(ccx, *callee, data, returnData) - }) + ccx.state.strategy.inner.clone().cheatcode_mock_call(ccx, *callee, data, returnData) } } @@ -85,9 +83,7 @@ impl Cheatcode for mockCalls_1Call { impl Cheatcode for mockCallRevert_0Call { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { callee, data, revertData } = self; - ccx.with_strategy(|strategy, ccx| { - strategy.cheatcode_mock_call_revert(ccx, *callee, data, revertData) - }) + ccx.state.strategy.inner.clone().cheatcode_mock_call_revert(ccx, *callee, data, revertData) } } diff --git a/crates/cheatcodes/src/fs.rs b/crates/cheatcodes/src/fs.rs index c70133a76..b99ed2945 100644 --- a/crates/cheatcodes/src/fs.rs +++ b/crates/cheatcodes/src/fs.rs @@ -283,7 +283,7 @@ impl Cheatcode for getArtifactPathByDeployedCodeCall { impl Cheatcode for getCodeCall { fn apply(&self, state: &mut Cheatcodes) -> Result { let Self { artifactPath: path } = self; - state.with_strategy(|strategy, state| strategy.get_artifact_code(state, path, false)) + state.strategy.inner.get_artifact_code(state, path, false) } } diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index af90c7a71..7ae10782c 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -7,7 +7,7 @@ use crate::{ DealRecord, GasRecord, }, script::{Broadcast, Wallets}, - strategy::CheatcodeInspectorStrategy, + strategy::Strategy, test::{ assume::AssumeNoRevert, expect::{self, ExpectedEmit, ExpectedRevert, ExpectedRevertKind}, @@ -69,7 +69,6 @@ pub use utils::CommonCreateInput; pub type Ecx<'a, 'b, 'c> = &'a mut EvmContext<&'b mut (dyn DatabaseExt + 'c)>; pub type InnerEcx<'a, 'b, 'c> = &'a mut InnerEvmContext<&'b mut (dyn DatabaseExt + 'c)>; -pub type Strategy<'a> = &'a mut dyn CheatcodeInspectorStrategy; /// Helper trait for obtaining complete [revm::Inspector] instance from mutable reference to /// [Cheatcodes]. @@ -528,7 +527,7 @@ pub struct Cheatcodes { pub wallets: Option, /// The behavior strategy. - pub strategy: Option>, + pub strategy: Strategy, } impl Clone for Cheatcodes { @@ -568,7 +567,7 @@ impl Clone for Cheatcodes { arbitrary_storage: self.arbitrary_storage.clone(), deprecated: self.deprecated.clone(), wallets: self.wallets.clone(), - strategy: self.strategy.as_ref().map(|s| s.new_cloned()), + strategy: self.strategy.clone(), } } } @@ -588,7 +587,7 @@ impl Cheatcodes { Self { fs_commit: true, labels: config.labels.clone(), - strategy: Some(config.strategy.clone()), + strategy: config.strategy.clone(), config, block: Default::default(), active_delegation: Default::default(), @@ -763,7 +762,8 @@ impl Cheatcodes { if ecx_inner.journaled_state.depth() == broadcast.depth { input.set_caller(broadcast.new_origin); - self.strategy.as_mut().unwrap().record_broadcastable_create_transactions( + self.strategy.inner.record_broadcastable_create_transactions( + self.strategy.context.as_mut(), self.config.clone(), &input, ecx_inner, @@ -801,9 +801,9 @@ impl Cheatcodes { }]); } - if let Some(result) = self.with_strategy(|strategy, cheatcodes| { - strategy.zksync_try_create(cheatcodes, ecx, &input, executor) - }) { + if let Some(result) = + self.strategy.inner.clone().zksync_try_create(self, ecx, &input, executor) + { return Some(result); } @@ -917,10 +917,7 @@ where { } } - self.strategy - .as_mut() - .expect("failed acquiring strategy") - .zksync_record_create_address(&outcome); + self.strategy.inner.zksync_record_create_address(self.strategy.context.as_mut(), &outcome); outcome } @@ -963,10 +960,12 @@ where { let prev = account.info.nonce; let nonce = prev.saturating_sub(1); account.info.nonce = nonce; - self.strategy - .as_mut() - .expect("failed acquiring strategy") - .zksync_sync_nonce(sender, nonce, ecx); + self.strategy.inner.zksync_sync_nonce( + self.strategy.context.as_mut(), + sender, + nonce, + ecx, + ); trace!(target: "cheatcodes", %sender, nonce, prev, "corrected nonce"); } @@ -1000,10 +999,7 @@ where { return None; } - self.strategy - .as_mut() - .expect("failed acquiring strategy") - .zksync_set_deployer_call_input(call); + self.strategy.inner.zksync_set_deployer_call_input(self.strategy.context.as_mut(), call); // Handle expected calls @@ -1137,17 +1133,15 @@ where { }) } - self.strategy - .as_mut() - .expect("failed acquiring strategy") - .record_broadcastable_call_transactions( - self.config.clone(), - call, - ecx_inner, - broadcast, - &mut self.broadcastable_transactions, - &mut self.active_delegation, - ); + self.strategy.inner.record_broadcastable_call_transactions( + self.strategy.context.as_mut(), + self.config.clone(), + call, + ecx_inner, + broadcast, + &mut self.broadcastable_transactions, + &mut self.active_delegation, + ); let account = ecx_inner.journaled_state.state().get_mut(&broadcast.new_origin).unwrap(); @@ -1220,9 +1214,8 @@ where { }]); } - if let Some(result) = self.with_strategy(|strategy, cheatcodes| { - strategy.zksync_try_call(cheatcodes, ecx, call, executor) - }) { + if let Some(result) = self.strategy.inner.clone().zksync_try_call(self, ecx, call, executor) + { return Some(result); } @@ -1264,17 +1257,6 @@ where { None => false, } } - - pub fn with_strategy(&mut self, mut f: F) -> R - where - F: FnMut(Strategy, &mut Self) -> R, - { - let mut strategy = self.strategy.take(); - let result = f(strategy.as_mut().expect("failed acquiring strategy").as_mut(), self); - self.strategy = strategy; - - result - } } impl Inspector<&mut dyn DatabaseExt> for Cheatcodes { @@ -1294,10 +1276,11 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes { self.gas_metering.paused_frames.push(interpreter.gas); } - self.strategy - .as_mut() - .expect("failed acquiring strategy") - .post_initialize_interp(interpreter, ecx); + self.strategy.inner.post_initialize_interp( + self.strategy.context.as_mut(), + interpreter, + ecx, + ); } #[inline] @@ -1342,8 +1325,7 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes { #[inline] fn step_end(&mut self, interpreter: &mut Interpreter, ecx: Ecx) { - if self.strategy.as_mut().expect("failed acquiring strategy").pre_step_end(interpreter, ecx) - { + if self.strategy.inner.pre_step_end(self.strategy.context.as_mut(), interpreter, ecx) { return; } diff --git a/crates/cheatcodes/src/lib.rs b/crates/cheatcodes/src/lib.rs index 56b6cfb43..33dec1c33 100644 --- a/crates/cheatcodes/src/lib.rs +++ b/crates/cheatcodes/src/lib.rs @@ -17,7 +17,6 @@ extern crate tracing; use alloy_primitives::Address; use foundry_evm_core::backend::DatabaseExt; -use inspector::Strategy; use revm::{ContextPrecompiles, InnerEvmContext}; use spec::Status; @@ -175,15 +174,4 @@ impl CheatsCtxt<'_, '_, '_, '_> { pub(crate) fn is_precompile(&self, address: &Address) -> bool { self.precompiles.contains(address) } - - pub(crate) fn with_strategy(&mut self, mut f: F) -> R - where - F: FnMut(Strategy, &mut Self) -> R, - { - let mut strategy = self.state.strategy.take(); - let result = f(strategy.as_mut().expect("failed acquiring strategy").as_mut(), self); - self.state.strategy = strategy; - - result - } } diff --git a/crates/cheatcodes/src/strategy.rs b/crates/cheatcodes/src/strategy.rs index 91154b660..6424b2f44 100644 --- a/crates/cheatcodes/src/strategy.rs +++ b/crates/cheatcodes/src/strategy.rs @@ -1,4 +1,4 @@ -use std::{fmt::Debug, sync::Arc}; +use std::{any::Any, fmt::Debug, sync::Arc}; use alloy_primitives::{Address, Bytes, FixedBytes, TxKind, U256}; use alloy_rpc_types::{TransactionInput, TransactionRequest}; @@ -18,39 +18,70 @@ use crate::{ CheatsConfig, CheatsCtxt, Result, }; +pub trait CheatcodeInspectorStrategyContext: Debug + Send + Sync + Any { + fn new_cloned(&self) -> Box; + fn as_any_mut(&mut self) -> &mut dyn Any; + fn as_any_ref(&self) -> &dyn Any; +} + +impl CheatcodeInspectorStrategyContext for () { + fn new_cloned(&self) -> Box { + Box::new(self.clone()) + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + fn as_any_ref(&self) -> &dyn Any { + self + } +} + +#[derive(Debug)] +pub struct Strategy { + pub inner: Box, + pub context: Box, +} + +pub fn new_evm_strategy() -> Strategy { + Strategy { inner: Box::new(EvmCheatcodeInspectorStrategy::default()), context: Box::new(()) } +} + +impl Clone for Strategy { + fn clone(&self) -> Self { + Self { inner: self.inner.new_cloned(), context: self.context.new_cloned() } + } +} + pub trait CheatcodeInspectorStrategy: Debug + Send + Sync + CheatcodeInspectorStrategyExt { fn name(&self) -> &'static str; fn new_cloned(&self) -> Box; /// Get nonce. - fn get_nonce(&mut self, ccx: &mut CheatsCtxt, address: Address) -> Result { + fn get_nonce(&self, ccx: &mut CheatsCtxt, address: Address) -> Result { let account = ccx.ecx.journaled_state.load_account(address, &mut ccx.ecx.db)?; Ok(account.info.nonce) } /// Called when the main test or script contract is deployed. - fn base_contract_deployed(&mut self) {} + fn base_contract_deployed(&self, _ctx: &mut dyn CheatcodeInspectorStrategyContext) {} /// Cheatcode: roll. - fn cheatcode_roll(&mut self, ccx: &mut CheatsCtxt, new_height: U256) -> Result { + fn cheatcode_roll(&self, ccx: &mut CheatsCtxt, new_height: U256) -> Result { ccx.ecx.env.block.number = new_height; Ok(Default::default()) } /// Cheatcode: warp. - fn cheatcode_warp(&mut self, ccx: &mut CheatsCtxt, new_timestamp: U256) -> Result { + fn cheatcode_warp(&self, ccx: &mut CheatsCtxt, new_timestamp: U256) -> Result { ccx.ecx.env.block.timestamp = new_timestamp; Ok(Default::default()) } /// Cheatcode: deal. - fn cheatcode_deal( - &mut self, - ccx: &mut CheatsCtxt, - address: Address, - new_balance: U256, - ) -> Result { + fn cheatcode_deal(&self, ccx: &mut CheatsCtxt, address: Address, new_balance: U256) -> Result { let account = journaled_account(ccx.ecx, address)?; let old_balance = std::mem::replace(&mut account.info.balance, new_balance); let record = DealRecord { address, old_balance, new_balance }; @@ -60,7 +91,7 @@ pub trait CheatcodeInspectorStrategy: Debug + Send + Sync + CheatcodeInspectorSt /// Cheatcode: etch. fn cheatcode_etch( - &mut self, + &self, ccx: &mut CheatsCtxt, target: Address, new_runtime_bytecode: &Bytes, @@ -73,12 +104,12 @@ pub trait CheatcodeInspectorStrategy: Debug + Send + Sync + CheatcodeInspectorSt } /// Cheatcode: getNonce. - fn cheatcode_get_nonce(&mut self, ccx: &mut CheatsCtxt, address: Address) -> Result { + fn cheatcode_get_nonce(&self, ccx: &mut CheatsCtxt, address: Address) -> Result { evm::get_nonce(ccx, &address) } /// Cheatcode: resetNonce. - fn cheatcode_reset_nonce(&mut self, ccx: &mut CheatsCtxt, account: Address) -> Result { + fn cheatcode_reset_nonce(&self, ccx: &mut CheatsCtxt, account: Address) -> Result { let account = journaled_account(ccx.ecx, account)?; // Per EIP-161, EOA nonces start at 0, but contract nonces // start at 1. Comparing by code_hash instead of code @@ -92,7 +123,7 @@ pub trait CheatcodeInspectorStrategy: Debug + Send + Sync + CheatcodeInspectorSt /// Cheatcode: setNonce. fn cheatcode_set_nonce( - &mut self, + &self, ccx: &mut CheatsCtxt, account: Address, new_nonce: u64, @@ -111,7 +142,7 @@ pub trait CheatcodeInspectorStrategy: Debug + Send + Sync + CheatcodeInspectorSt /// Cheatcode: setNonceUnsafe. fn cheatcode_set_nonce_unsafe( - &mut self, + &self, ccx: &mut CheatsCtxt, account: Address, new_nonce: u64, @@ -123,7 +154,7 @@ pub trait CheatcodeInspectorStrategy: Debug + Send + Sync + CheatcodeInspectorSt /// Mocks a call to return with a value. fn cheatcode_mock_call( - &mut self, + &self, ccx: &mut CheatsCtxt, callee: Address, data: &Bytes, @@ -136,7 +167,7 @@ pub trait CheatcodeInspectorStrategy: Debug + Send + Sync + CheatcodeInspectorSt /// Mocks a call to revert with a value. fn cheatcode_mock_call_revert( - &mut self, + &self, ccx: &mut CheatsCtxt, callee: Address, data: &Bytes, @@ -154,7 +185,8 @@ pub trait CheatcodeInspectorStrategy: Debug + Send + Sync + CheatcodeInspectorSt /// Record broadcastable transaction during CREATE. fn record_broadcastable_create_transactions( - &mut self, + &self, + _ctx: &mut dyn CheatcodeInspectorStrategyContext, config: Arc, input: &dyn CommonCreateInput, ecx_inner: InnerEcx, @@ -164,7 +196,8 @@ pub trait CheatcodeInspectorStrategy: Debug + Send + Sync + CheatcodeInspectorSt /// Record broadcastable transaction during CALL. fn record_broadcastable_call_transactions( - &mut self, + &self, + _ctx: &mut dyn CheatcodeInspectorStrategyContext, config: Arc, input: &CallInputs, ecx_inner: InnerEcx, @@ -173,35 +206,55 @@ pub trait CheatcodeInspectorStrategy: Debug + Send + Sync + CheatcodeInspectorSt active_delegation: &mut Option, ); - fn post_initialize_interp(&mut self, _interpreter: &mut Interpreter, _ecx: Ecx) {} + fn post_initialize_interp( + &self, + _ctx: &mut dyn CheatcodeInspectorStrategyContext, + _interpreter: &mut Interpreter, + _ecx: Ecx, + ) { + } /// Used to override opcode behaviors. Returns true if handled. - fn pre_step_end(&mut self, _interpreter: &mut Interpreter, _ecx: Ecx) -> bool { + fn pre_step_end( + &self, + _ctx: &mut dyn CheatcodeInspectorStrategyContext, + _interpreter: &mut Interpreter, + _ecx: Ecx, + ) -> bool { false } } /// We define this in our fork pub trait CheatcodeInspectorStrategyExt { - fn zksync_cheatcode_skip_zkvm(&mut self) -> Result { + fn zksync_cheatcode_skip_zkvm( + &self, + _ctx: &mut dyn CheatcodeInspectorStrategyContext, + ) -> Result { Ok(Default::default()) } fn zksync_cheatcode_set_paymaster( - &mut self, + &self, + _ctx: &mut dyn CheatcodeInspectorStrategyContext, _paymaster_address: Address, _paymaster_input: &Bytes, ) -> Result { Ok(Default::default()) } - fn zksync_cheatcode_use_factory_deps(&mut self, _name: String) -> Result { + fn zksync_cheatcode_use_factory_deps( + &self, + _ctx: &mut dyn CheatcodeInspectorStrategyContext, + _name: String, + ) -> Result { Ok(Default::default()) } #[allow(clippy::too_many_arguments)] fn zksync_cheatcode_register_contract( - &mut self, + &self, + _ctx: &mut dyn CheatcodeInspectorStrategyContext, _name: String, _zk_bytecode_hash: FixedBytes<32>, _zk_deployed_bytecode: Vec, @@ -213,16 +266,39 @@ pub trait CheatcodeInspectorStrategyExt { Ok(Default::default()) } - fn zksync_cheatcode_select_zk_vm(&mut self, _data: InnerEcx, _enable: bool) {} + fn zksync_cheatcode_select_zk_vm( + &self, + _ctx: &mut dyn CheatcodeInspectorStrategyContext, + _data: InnerEcx, + _enable: bool, + ) { + } - fn zksync_record_create_address(&mut self, _outcome: &CreateOutcome) {} + fn zksync_record_create_address( + &self, + _ctx: &mut dyn CheatcodeInspectorStrategyContext, + _outcome: &CreateOutcome, + ) { + } - fn zksync_sync_nonce(&mut self, _sender: Address, _nonce: u64, _ecx: Ecx) {} + fn zksync_sync_nonce( + &self, + _ctx: &mut dyn CheatcodeInspectorStrategyContext, + _sender: Address, + _nonce: u64, + _ecx: Ecx, + ) { + } - fn zksync_set_deployer_call_input(&mut self, _call: &mut CallInputs) {} + fn zksync_set_deployer_call_input( + &self, + _ctx: &mut dyn CheatcodeInspectorStrategyContext, + _call: &mut CallInputs, + ) { + } fn zksync_try_create( - &mut self, + &self, _state: &mut Cheatcodes, _ecx: Ecx<'_, '_, '_>, _input: &dyn CommonCreateInput, @@ -232,7 +308,7 @@ pub trait CheatcodeInspectorStrategyExt { } fn zksync_try_call( - &mut self, + &self, _state: &mut Cheatcodes, _ecx: Ecx, _input: &CallInputs, @@ -241,7 +317,13 @@ pub trait CheatcodeInspectorStrategyExt { None } - fn zksync_select_fork_vm(&mut self, _data: InnerEcx, _fork_id: LocalForkId) {} + fn zksync_select_fork_vm( + &self, + _ctx: &mut dyn CheatcodeInspectorStrategyContext, + _data: InnerEcx, + _fork_id: LocalForkId, + ) { + } } #[derive(Debug, Default, Clone)] @@ -257,7 +339,8 @@ impl CheatcodeInspectorStrategy for EvmCheatcodeInspectorStrategy { } fn record_broadcastable_create_transactions( - &mut self, + &self, + _ctx: &mut dyn CheatcodeInspectorStrategyContext, _config: Arc, input: &dyn CommonCreateInput, ecx_inner: InnerEcx, @@ -288,7 +371,8 @@ impl CheatcodeInspectorStrategy for EvmCheatcodeInspectorStrategy { } fn record_broadcastable_call_transactions( - &mut self, + &self, + _ctx: &mut dyn CheatcodeInspectorStrategyContext, _config: Arc, call: &CallInputs, ecx_inner: InnerEcx, diff --git a/crates/cheatcodes/src/test.rs b/crates/cheatcodes/src/test.rs index de3f1eb72..9c2e811b3 100644 --- a/crates/cheatcodes/src/test.rs +++ b/crates/cheatcodes/src/test.rs @@ -15,7 +15,11 @@ impl Cheatcode for zkVmCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { enable } = *self; - ccx.with_strategy(|strategy, ccx| strategy.zksync_cheatcode_select_zk_vm(ccx.ecx, enable)); + ccx.state.strategy.inner.zksync_cheatcode_select_zk_vm( + ccx.state.strategy.context.as_mut(), + ccx.ecx, + enable, + ); Ok(Default::default()) } @@ -23,23 +27,28 @@ impl Cheatcode for zkVmCall { impl Cheatcode for zkVmSkipCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { - ccx.with_strategy(|strategy, _ccx| strategy.zksync_cheatcode_skip_zkvm()) + ccx.state.strategy.inner.zksync_cheatcode_skip_zkvm(ccx.state.strategy.context.as_mut()) } } impl Cheatcode for zkUsePaymasterCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { paymaster_address, paymaster_input } = self; - ccx.with_strategy(|strategy, _ccx| { - strategy.zksync_cheatcode_set_paymaster(*paymaster_address, paymaster_input) - }) + ccx.state.strategy.inner.zksync_cheatcode_set_paymaster( + ccx.state.strategy.context.as_mut(), + *paymaster_address, + paymaster_input, + ) } } impl Cheatcode for zkUseFactoryDepCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { name } = self; - ccx.with_strategy(|strategy, _ccx| strategy.zksync_cheatcode_use_factory_deps(name.clone())) + ccx.state + .strategy + .inner + .zksync_cheatcode_use_factory_deps(ccx.state.strategy.context.as_mut(), name.clone()) } } @@ -54,17 +63,16 @@ impl Cheatcode for zkRegisterContractCall { zkDeployedBytecode, } = self; - ccx.with_strategy(|strategy, _ccx| { - 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(), - ) - }) + ccx.state.strategy.inner.zksync_cheatcode_register_contract( + ccx.state.strategy.context.as_mut(), + 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 7a72c111e..285d55496 100644 --- a/crates/chisel/src/executor.rs +++ b/crates/chisel/src/executor.rs @@ -316,14 +316,14 @@ impl SessionSource { let env = self.config.evm_opts.evm_env().await.expect("Could not instantiate fork environment"); - let strategy = utils::get_executor_strategy(&self.config.foundry_config); + let mut strategy = utils::get_executor_strategy(&self.config.foundry_config); // Create an in-memory backend let backend = match self.config.backend.take() { Some(backend) => backend, None => { let fork = self.config.evm_opts.get_fork(&self.config.foundry_config, env.clone()); - let backend = Backend::spawn(fork, strategy.new_backend_strategy()); + let backend = Backend::spawn(fork, strategy.inner.new_backend_strategy()); self.config.backend = Some(backend.clone()); backend } @@ -339,7 +339,7 @@ impl SessionSource { None, None, Some(self.solc.version.clone()), - strategy.new_cheatcode_inspector_strategy(), + strategy.inner.new_cheatcode_inspector_strategy(strategy.context.as_mut()), ) .into(), ) diff --git a/crates/cli/src/utils/mod.rs b/crates/cli/src/utils/mod.rs index 48cf75cce..e06ff934a 100644 --- a/crates/cli/src/utils/mod.rs +++ b/crates/cli/src/utils/mod.rs @@ -9,8 +9,8 @@ use foundry_common::{ shell, }; use foundry_config::{Chain, Config}; -use foundry_evm::executors::strategy::{EvmExecutorStrategy, ExecutorStrategy}; -use foundry_strategy_zksync::ZksyncExecutorStrategy; +use foundry_evm::executors::strategy::{new_evm_strategy, Strategy}; +use foundry_strategy_zksync::new_zkysnc_strategy; use serde::de::DeserializeOwned; use std::{ ffi::OsStr, @@ -93,13 +93,13 @@ pub fn get_provider(config: &Config) -> Result { get_provider_builder(config)?.build() } -pub fn get_executor_strategy(config: &Config) -> Box { +pub fn get_executor_strategy(config: &Config) -> Strategy { if config.zksync.should_compile() { info!("using zksync strategy"); - Box::new(ZksyncExecutorStrategy::default()) + new_zkysnc_strategy() } else { info!("using evm strategy"); - Box::new(EvmExecutorStrategy::default()) + new_evm_strategy() } } diff --git a/crates/evm/core/src/backend/cow.rs b/crates/evm/core/src/backend/cow.rs index f95b9087e..f53b6ce15 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::{strategy::BackendStrategy, BackendError, ForkInfo}; +use super::{strategy::Strategy, BackendError, ForkInfo}; use crate::{ backend::{ diagnostic::RevertDiagnostic, Backend, DatabaseExt, LocalForkId, RevertStateSnapshotAction, @@ -92,8 +92,8 @@ impl DatabaseExt for CowBackend<'_> { self.backend.to_mut().get_fork_info(id) } - fn get_strategy(&mut self) -> &mut dyn BackendStrategy { - self.backend.to_mut().strategy.as_mut() + fn get_strategy(&mut self) -> &mut Strategy { + &mut self.backend.to_mut().strategy } 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 fbd5bf0a2..4078e7736 100644 --- a/crates/evm/core/src/backend/mod.rs +++ b/crates/evm/core/src/backend/mod.rs @@ -30,7 +30,7 @@ use std::{ collections::{BTreeMap, HashSet}, time::Instant, }; -use strategy::{BackendStrategy, BackendStrategyForkInfo}; +use strategy::{BackendStrategyForkInfo, Strategy}; mod diagnostic; pub use diagnostic::RevertDiagnostic; @@ -101,8 +101,7 @@ pub trait DatabaseExt: Database + DatabaseCommit { /// and the the fork environment. fn get_fork_info(&mut self, id: LocalForkId) -> eyre::Result; - /// Retrieve the strategy. - fn get_strategy(&mut self) -> &mut dyn BackendStrategy; + fn get_strategy(&mut self) -> &mut Strategy; /// Reverts the snapshot if it exists /// @@ -460,7 +459,7 @@ struct _ObjectSafe(dyn DatabaseExt); #[must_use] pub struct Backend { /// The behavior strategy. - pub strategy: Box, + pub strategy: Strategy, /// The access point for managing forks forks: MultiFork, @@ -496,13 +495,13 @@ pub struct Backend { impl Clone for Backend { fn clone(&self) -> Self { Self { - strategy: self.strategy.new_cloned(), forks: self.forks.clone(), mem_db: self.mem_db.clone(), fork_init_journaled_state: self.fork_init_journaled_state.clone(), active_fork_ids: self.active_fork_ids, inner: self.inner.clone(), fork_url_type: self.fork_url_type.clone(), + strategy: self.strategy.clone(), } } } @@ -512,7 +511,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: Box) -> Self { + pub fn spawn(fork: Option, strategy: Strategy) -> Self { Self::new(MultiFork::spawn(), fork, strategy) } @@ -522,11 +521,7 @@ impl Backend { /// database. /// /// Prefer using [`spawn`](Self::spawn) instead. - pub fn new( - forks: MultiFork, - fork: Option, - strategy: Box, - ) -> Self { + pub fn new(forks: MultiFork, fork: Option, strategy: Strategy) -> Self { trace!(target: "backend", forking_mode=?fork.is_some(), "creating executor backend"); // Note: this will take of registering the `fork` let inner = BackendInner { @@ -565,10 +560,10 @@ impl Backend { /// Creates a new instance of `Backend` with fork added to the fork database and sets the fork /// as active pub(crate) fn new_with_fork( + strategy: Strategy, id: &ForkId, fork: Fork, journaled_state: JournaledState, - strategy: Box, ) -> Self { let mut backend = Self::spawn(None, strategy); let fork_ids = backend.inner.insert_new_fork(id.clone(), fork.db, journaled_state); @@ -586,7 +581,7 @@ impl Backend { active_fork_ids: None, inner: Default::default(), fork_url_type: Default::default(), - strategy: self.strategy.new_cloned(), + strategy: self.strategy.clone(), } } @@ -907,6 +902,7 @@ impl Backend { trace!(tx=?tx.tx_hash(), "committing transaction"); commit_transaction( + &mut self.strategy, &tx.inner, env.clone(), journaled_state, @@ -914,7 +910,6 @@ impl Backend { &fork_id, &persistent_accounts, &mut NoOpInspector, - self.strategy.as_mut(), )?; } @@ -938,8 +933,8 @@ impl DatabaseExt for Backend { Ok(ForkInfo { fork_type, fork_env }) } - fn get_strategy(&mut self) -> &mut dyn BackendStrategy { - self.strategy.as_mut() + fn get_strategy(&mut self) -> &mut Strategy { + &mut self.strategy } fn snapshot_state(&mut self, journaled_state: &JournaledState, env: &Env) -> U256 { @@ -1169,9 +1164,12 @@ impl DatabaseExt for Backend { caller_account.into() }); - self.strategy.update_fork_db( + let active_fork = self.active_fork_ids.map(|(_, idx)| self.inner.get_fork(idx)); + // let active_fork = self.active_fork().cloned(); + self.strategy.inner.update_fork_db( + self.strategy.context.as_mut(), BackendStrategyForkInfo { - active_fork: self.active_fork(), + active_fork, active_type: current_fork_type, target_type: target_fork_type, }, @@ -1206,7 +1204,7 @@ impl DatabaseExt for Backend { let (fork_id, backend, fork_env) = self.forks.roll_fork(self.inner.ensure_fork_id(id).cloned()?, block_number)?; // this will update the local mapping - self.inner.roll_fork(id, fork_id, backend, self.strategy.as_mut())?; + self.inner.roll_fork(&mut self.strategy, id, fork_id, backend)?; if let Some((active_id, active_idx)) = self.active_fork_ids { // the currently active fork is the targeted fork of this call @@ -1228,7 +1226,8 @@ impl DatabaseExt for Backend { active.journaled_state.depth = journaled_state.depth; for addr in persistent_addrs { - self.strategy.merge_journaled_state_data( + self.strategy.inner.merge_journaled_state_data( + self.strategy.context.as_mut(), addr, journaled_state, &mut active.journaled_state, @@ -1245,7 +1244,8 @@ impl DatabaseExt for Backend { for (addr, acc) in journaled_state.state.iter() { if acc.is_created() { if acc.is_touched() { - self.strategy.merge_journaled_state_data( + self.strategy.inner.merge_journaled_state_data( + self.strategy.context.as_mut(), *addr, journaled_state, &mut active.journaled_state, @@ -1264,6 +1264,7 @@ impl DatabaseExt for Backend { fn roll_fork_to_transaction( &mut self, + id: Option, transaction: B256, env: &mut Env, @@ -1319,6 +1320,7 @@ impl DatabaseExt for Backend { let env = self.env_with_handler_cfg(env); let fork = self.inner.get_fork_by_id_mut(id)?; commit_transaction( + &mut self.strategy, &tx, env, journaled_state, @@ -1326,7 +1328,6 @@ impl DatabaseExt for Backend { &fork_id, &persistent_accounts, inspector, - self.strategy.as_mut(), ) } @@ -1805,10 +1806,10 @@ impl BackendInner { pub fn roll_fork( &mut self, + strategy: &mut Strategy, id: LocalForkId, new_fork_id: ForkId, backend: SharedBackend, - strategy: &mut dyn BackendStrategy, ) -> eyre::Result { let fork_id = self.ensure_fork_id(id)?; let idx = self.ensure_fork_index(fork_id)?; @@ -1817,7 +1818,12 @@ impl BackendInner { // we initialize a _new_ `ForkDB` but keep the state of persistent accounts let mut new_db = ForkDB::new(backend); for addr in self.persistent_accounts.iter().copied() { - strategy.merge_db_account_data(addr, &active.db, &mut new_db); + strategy.inner.merge_db_account_data( + strategy.context.as_mut(), + addr, + &active.db, + &mut new_db, + ); } active.db = new_db; } @@ -1930,6 +1936,7 @@ fn update_env_block(env: &mut Env, block: &AnyRpcBlock) { /// state, with an inspector. #[allow(clippy::too_many_arguments)] fn commit_transaction( + strategy: &mut Strategy, tx: &Transaction, mut env: EnvWithHandlerCfg, journaled_state: &mut JournaledState, @@ -1937,7 +1944,6 @@ fn commit_transaction( fork_id: &ForkId, persistent_accounts: &HashSet
, inspector: &mut dyn InspectorExt, - strategy: &mut dyn BackendStrategy, ) -> 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 @@ -1952,7 +1958,7 @@ fn commit_transaction( let fork = fork.clone(); let journaled_state = journaled_state.clone(); let depth = journaled_state.depth; - let mut db = Backend::new_with_fork(fork_id, fork, journaled_state, strategy.new_cloned()); + let mut db = Backend::new_with_fork(strategy.clone(), fork_id, fork, journaled_state); let mut evm = crate::utils::new_evm_with_inspector(&mut db as _, env, inspector); // Adjust inner EVM depth to ensure that inspectors receive accurate data. @@ -2005,7 +2011,7 @@ fn apply_state_changeset( #[allow(clippy::needless_return)] mod tests { use crate::{ - backend::{strategy::EvmBackendStrategy, Backend}, + backend::{strategy::new_evm_strategy, Backend}, fork::CreateFork, opts::EvmOpts, }; @@ -2039,7 +2045,7 @@ mod tests { evm_opts, }; - let backend = Backend::spawn(Some(fork), Box::new(EvmBackendStrategy)); + let backend = Backend::spawn(Some(fork), new_evm_strategy()); // some rng contract from etherscan let address: Address = "63091244180ae240c87d1f528f5f269134cb07b3".parse().unwrap(); diff --git a/crates/evm/core/src/backend/strategy.rs b/crates/evm/core/src/backend/strategy.rs index 18ebe9e6b..74dd3023a 100644 --- a/crates/evm/core/src/backend/strategy.rs +++ b/crates/evm/core/src/backend/strategy.rs @@ -1,4 +1,4 @@ -use std::fmt::Debug; +use std::{any::Any, fmt::Debug}; use super::{BackendInner, Fork, ForkDB, ForkType, FoundryEvmInMemoryDB}; use alloy_primitives::{Address, U256}; @@ -11,6 +11,42 @@ pub struct BackendStrategyForkInfo<'a> { pub target_type: ForkType, } +pub trait BackendStrategyContext: Debug + Send + Sync + Any { + fn new_cloned(&self) -> Box; + fn as_any_ref(&self) -> &dyn Any; + fn as_any_mut(&mut self) -> &mut dyn Any; +} + +impl BackendStrategyContext for () { + fn new_cloned(&self) -> Box { + Box::new(self.clone()) + } + + fn as_any_ref(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + +#[derive(Debug)] +pub struct Strategy { + pub inner: Box, + pub context: Box, +} + +pub fn new_evm_strategy() -> Strategy { + Strategy { inner: Box::new(EvmBackendStrategy), context: Box::new(()) } +} + +impl Clone for Strategy { + fn clone(&self) -> Self { + Self { inner: self.inner.new_cloned(), context: self.context.new_cloned() } + } +} + pub trait BackendStrategy: Debug + Send + Sync + BackendStrategyExt { fn name(&self) -> &'static str; @@ -19,6 +55,7 @@ pub trait BackendStrategy: Debug + Send + Sync + BackendStrategyExt { /// When creating or switching forks, we update the AccountInfo of the contract fn update_fork_db( &self, + ctx: &mut dyn BackendStrategyContext, fork_info: BackendStrategyForkInfo<'_>, mem_db: &FoundryEvmInMemoryDB, backend_inner: &BackendInner, @@ -29,12 +66,19 @@ pub trait BackendStrategy: Debug + Send + Sync + BackendStrategyExt { /// Clones the account data from the `active_journaled_state` into the `fork_journaled_state` fn merge_journaled_state_data( &self, + ctx: &mut dyn BackendStrategyContext, addr: Address, active_journaled_state: &JournaledState, fork_journaled_state: &mut JournaledState, ); - fn merge_db_account_data(&self, addr: Address, active: &ForkDB, fork_db: &mut ForkDB); + fn merge_db_account_data( + &self, + ctx: &mut dyn BackendStrategyContext, + addr: Address, + active: &ForkDB, + fork_db: &mut ForkDB, + ); } pub trait BackendStrategyExt { @@ -44,7 +88,13 @@ pub trait BackendStrategyExt { /// 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) {} + fn zksync_save_immutable_storage( + &self, + _ctx: &mut dyn BackendStrategyContext, + _addr: Address, + _keys: HashSet, + ) { + } } struct _ObjectSafe(dyn BackendStrategy); @@ -63,6 +113,7 @@ impl BackendStrategy for EvmBackendStrategy { fn update_fork_db( &self, + _ctx: &mut dyn BackendStrategyContext, fork_info: BackendStrategyForkInfo<'_>, mem_db: &FoundryEvmInMemoryDB, backend_inner: &BackendInner, @@ -80,6 +131,7 @@ impl BackendStrategy for EvmBackendStrategy { fn merge_journaled_state_data( &self, + _ctx: &mut dyn BackendStrategyContext, addr: Address, active_journaled_state: &JournaledState, fork_journaled_state: &mut JournaledState, @@ -91,7 +143,13 @@ impl BackendStrategy for EvmBackendStrategy { ); } - fn merge_db_account_data(&self, addr: Address, active: &ForkDB, fork_db: &mut ForkDB) { + fn merge_db_account_data( + &self, + _ctx: &mut dyn BackendStrategyContext, + addr: Address, + active: &ForkDB, + fork_db: &mut ForkDB, + ) { EvmBackendMergeStrategy::merge_db_account_data(addr, active, fork_db); } } @@ -199,3 +257,9 @@ impl EvmBackendMergeStrategy { fork_db.accounts.insert(addr, acc); } } + +impl Clone for Box { + fn clone(&self) -> Self { + self.new_cloned() + } +} diff --git a/crates/evm/evm/src/executors/builder.rs b/crates/evm/evm/src/executors/builder.rs index 40fc2fa31..3dc7009df 100644 --- a/crates/evm/evm/src/executors/builder.rs +++ b/crates/evm/evm/src/executors/builder.rs @@ -2,7 +2,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::Strategy; /// The builder that allows to configure an evm [`Executor`] which a stack of optional /// [`revm::Inspector`]s, such as [`Cheatcodes`]. @@ -76,7 +76,7 @@ impl ExecutorBuilder { /// Builds the executor as configured. #[inline] - pub fn build(self, env: Env, db: Backend, strategy: Box) -> Executor { + pub fn build(self, env: Env, db: Backend, strategy: Strategy) -> Executor { let Self { mut stack, gas_limit, spec_id, legacy_assertions } = self; if stack.block.is_none() { stack.block = Some(env.block.clone()); diff --git a/crates/evm/evm/src/executors/mod.rs b/crates/evm/evm/src/executors/mod.rs index 1e856aa28..352148fc2 100644 --- a/crates/evm/evm/src/executors/mod.rs +++ b/crates/evm/evm/src/executors/mod.rs @@ -37,7 +37,7 @@ use revm::{ }, }; use std::borrow::Cow; -use strategy::ExecutorStrategy; +use strategy::Strategy; mod builder; pub use builder::ExecutorBuilder; @@ -93,7 +93,7 @@ pub struct Executor { /// Whether `failed()` should be called on the test contract to determine if the test failed. legacy_assertions: bool, - strategy: Option>, + strategy: Strategy, } impl Clone for Executor { @@ -104,7 +104,7 @@ impl Clone for Executor { inspector: self.inspector.clone(), gas_limit: self.gas_limit, legacy_assertions: self.legacy_assertions, - strategy: self.strategy.as_ref().map(|s| s.new_cloned()), + strategy: self.strategy.clone(), } } } @@ -124,7 +124,7 @@ impl Executor { inspector: InspectorStack, gas_limit: u64, legacy_assertions: bool, - strategy: Box, + strategy: Strategy, ) -> Self { // Need to create a non-empty contract on the cheatcodes address so `extcodesize` checks // do not fail. @@ -139,7 +139,7 @@ impl Executor { }, ); - Self { backend, env, inspector, gas_limit, legacy_assertions, strategy: Some(strategy) } + Self { backend, env, inspector, gas_limit, legacy_assertions, strategy } } fn clone_with_backend(&self, backend: Backend) -> Self { @@ -150,7 +150,7 @@ impl Executor { self.inspector().clone(), self.gas_limit, self.legacy_assertions, - self.strategy.as_ref().map(|s| s.new_cloned()).expect("failed acquiring strategy"), + self.strategy.clone(), ) } @@ -214,21 +214,10 @@ impl Executor { Ok(()) } - pub fn with_strategy(&mut self, mut f: F) -> R - where - F: FnMut(&mut dyn ExecutorStrategy, &mut Self) -> R, - { - let mut strategy = self.strategy.take(); - let result = f(strategy.as_mut().expect("failed acquiring strategy").as_mut(), self); - self.strategy = strategy; - - result - } - /// Set the balance of an account. pub fn set_balance(&mut self, address: Address, amount: U256) -> BackendResult<()> { trace!(?address, ?amount, "setting account balance"); - self.with_strategy(|strategy, executor| strategy.set_balance(executor, address, amount)) + self.strategy.inner.clone().set_balance(self, address, amount) } /// Gets the balance of an account @@ -238,7 +227,7 @@ impl Executor { /// Set the nonce of an account. pub fn set_nonce(&mut self, address: Address, nonce: u64) -> BackendResult<()> { - self.with_strategy(|strategy, executor| strategy.set_nonce(executor, address, nonce)) + self.strategy.inner.clone().set_nonce(self, address, nonce) } /// Returns the nonce of an account. @@ -271,10 +260,7 @@ impl Executor { #[inline] pub fn set_transaction_other_fields(&mut self, other_fields: OtherFields) { - self.strategy - .as_mut() - .expect("failed acquiring strategy") - .set_inspect_context(other_fields); + self.strategy.inner.set_inspect_context(self.strategy.context.as_mut(), other_fields); } /// Deploys a contract and commits the new state to the underlying database. @@ -457,12 +443,12 @@ impl Executor { backend.is_initialized = false; backend.spec_id = env.spec_id(); - let result = self - .strategy - .as_ref() - .expect("failed acquiring strategy") - .new_cloned() - .call_inspect(&mut backend, &mut env, &mut inspector)?; + let result = self.strategy.inner.call_inspect( + self.strategy.context.as_ref(), + &mut backend, + &mut env, + &mut inspector, + )?; convert_executed_result( env.clone(), @@ -475,24 +461,27 @@ impl Executor { /// Execute the transaction configured in `env.tx`. #[instrument(name = "transact", level = "debug", skip_all)] pub fn transact_with_env(&mut self, mut env: EnvWithHandlerCfg) -> eyre::Result { - self.with_strategy(|strategy, executor| { - let mut inspector = executor.inspector.clone(); - let backend = &mut executor.backend; - backend.initialize(&env); - - let result_and_state = - strategy.transact_inspect(backend, &mut env, &executor.env, &mut inspector)?; - - let mut result = convert_executed_result( - env.clone(), - inspector, - result_and_state, - backend.has_state_snapshot_failure(), - )?; - - executor.commit(&mut result); - Ok(result) - }) + let mut inspector = self.inspector.clone(); + let backend = &mut self.backend; + backend.initialize(&env); + + let result_and_state = self.strategy.inner.transact_inspect( + self.strategy.context.as_mut(), + backend, + &mut env, + &self.env, + &mut inspector, + )?; + + let mut result = convert_executed_result( + env.clone(), + inspector, + result_and_state, + backend.has_state_snapshot_failure(), + )?; + + self.commit(&mut result); + Ok(result) } /// Commit the changeset to the database and adjust `self.inspector_config` values according to diff --git a/crates/evm/evm/src/executors/strategy.rs b/crates/evm/evm/src/executors/strategy.rs index 06919d8ae..a2cc03cd2 100644 --- a/crates/evm/evm/src/executors/strategy.rs +++ b/crates/evm/evm/src/executors/strategy.rs @@ -1,14 +1,11 @@ -use std::fmt::Debug; +use std::{any::Any, fmt::Debug}; use alloy_primitives::{Address, U256}; use alloy_serde::OtherFields; use eyre::{Context, Result}; -use foundry_cheatcodes::strategy::{CheatcodeInspectorStrategy, EvmCheatcodeInspectorStrategy}; +use foundry_cheatcodes::strategy::EvmCheatcodeInspectorStrategy; use foundry_evm_core::{ - backend::{ - strategy::{BackendStrategy, EvmBackendStrategy}, - BackendResult, DatabaseExt, - }, + backend::{BackendResult, DatabaseExt}, InspectorExt, }; use foundry_zksync_compiler::DualCompiledContracts; @@ -19,56 +16,99 @@ use revm::{ use super::Executor; +pub trait ExecutorStrategyContext: Debug + Send + Sync + Any { + fn new_cloned(&self) -> Box; + fn as_any_ref(&self) -> &dyn Any; + fn as_any_mut(&mut self) -> &mut dyn Any; +} + +impl ExecutorStrategyContext for () { + fn new_cloned(&self) -> Box { + Box::new(self.clone()) + } + + fn as_any_ref(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + +#[derive(Debug)] +pub struct Strategy { + pub inner: Box, + pub context: Box, +} + +pub fn new_evm_strategy() -> Strategy { + Strategy { inner: Box::new(EvmExecutorStrategy::default()), context: Box::new(()) } +} + +impl Clone for Strategy { + fn clone(&self) -> Self { + Self { inner: self.inner.new_cloned(), context: self.context.new_cloned() } + } +} + pub trait ExecutorStrategy: Debug + Send + Sync + ExecutorStrategyExt { fn name(&self) -> &'static str; fn new_cloned(&self) -> Box; fn set_balance( - &mut self, + &self, executor: &mut Executor, address: Address, amount: U256, ) -> BackendResult<()>; - fn set_nonce( - &mut self, - executor: &mut Executor, - address: Address, - nonce: u64, - ) -> BackendResult<()>; + fn set_nonce(&self, executor: &mut Executor, address: Address, nonce: u64) + -> BackendResult<()>; - fn set_inspect_context(&mut self, other_fields: OtherFields); + fn set_inspect_context(&self, ctx: &mut dyn ExecutorStrategyContext, other_fields: OtherFields); fn call_inspect( &self, + ctx: &dyn ExecutorStrategyContext, db: &mut dyn DatabaseExt, env: &mut EnvWithHandlerCfg, inspector: &mut dyn InspectorExt, ) -> eyre::Result; fn transact_inspect( - &mut self, + &self, + ctx: &mut dyn ExecutorStrategyContext, db: &mut dyn DatabaseExt, env: &mut EnvWithHandlerCfg, _executor_env: &EnvWithHandlerCfg, inspector: &mut dyn InspectorExt, ) -> eyre::Result; - fn new_backend_strategy(&self) -> Box; - fn new_cheatcode_inspector_strategy(&self) -> Box; + fn new_backend_strategy(&self) -> foundry_evm_core::backend::strategy::Strategy; + fn new_cheatcode_inspector_strategy( + &self, + ctx: &dyn ExecutorStrategyContext, + ) -> foundry_cheatcodes::strategy::Strategy; // TODO perhaps need to create fresh strategies as well } pub trait ExecutorStrategyExt { fn zksync_set_dual_compiled_contracts( - &mut self, + &self, + _ctx: &mut dyn ExecutorStrategyContext, _dual_compiled_contracts: DualCompiledContracts, ) { } - fn zksync_set_fork_env(&mut self, _fork_url: &str, _env: &Env) -> Result<()> { + fn zksync_set_fork_env( + &self, + _ctx: &mut dyn ExecutorStrategyContext, + _fork_url: &str, + _env: &Env, + ) -> Result<()> { Ok(()) } } @@ -85,7 +125,12 @@ impl ExecutorStrategy for EvmExecutorStrategy { Box::new(self.clone()) } - fn set_inspect_context(&mut self, _other_fields: OtherFields) {} + fn set_inspect_context( + &self, + _ctx: &mut dyn ExecutorStrategyContext, + _other_fields: OtherFields, + ) { + } /// Executes the configured test call of the `env` without committing state changes. /// @@ -93,6 +138,7 @@ impl ExecutorStrategy for EvmExecutorStrategy { /// update the given `env` with the new values. fn call_inspect( &self, + _ctx: &dyn ExecutorStrategyContext, db: &mut dyn DatabaseExt, env: &mut EnvWithHandlerCfg, inspector: &mut dyn InspectorExt, @@ -112,7 +158,8 @@ impl ExecutorStrategy for EvmExecutorStrategy { /// Note: in case there are any cheatcodes executed that modify the environment, this will /// update the given `env` with the new values. fn transact_inspect( - &mut self, + &self, + _ctx: &mut dyn ExecutorStrategyContext, db: &mut dyn DatabaseExt, env: &mut EnvWithHandlerCfg, _executor_env: &EnvWithHandlerCfg, @@ -128,7 +175,7 @@ impl ExecutorStrategy for EvmExecutorStrategy { } fn set_balance( - &mut self, + &self, executor: &mut Executor, address: Address, amount: U256, @@ -142,7 +189,7 @@ impl ExecutorStrategy for EvmExecutorStrategy { } fn set_nonce( - &mut self, + &self, executor: &mut Executor, address: Address, nonce: u64, @@ -154,13 +201,25 @@ impl ExecutorStrategy for EvmExecutorStrategy { Ok(()) } - fn new_backend_strategy(&self) -> Box { - Box::new(EvmBackendStrategy) + fn new_backend_strategy(&self) -> foundry_evm_core::backend::strategy::Strategy { + foundry_evm_core::backend::strategy::new_evm_strategy() } - fn new_cheatcode_inspector_strategy(&self) -> Box { - Box::new(EvmCheatcodeInspectorStrategy::default()) + fn new_cheatcode_inspector_strategy( + &self, + _ctx: &dyn ExecutorStrategyContext, + ) -> foundry_cheatcodes::strategy::Strategy { + foundry_cheatcodes::strategy::Strategy { + inner: Box::new(EvmCheatcodeInspectorStrategy::default()), + context: Box::new(()), + } } } impl ExecutorStrategyExt for EvmExecutorStrategy {} + +impl Clone for Box { + fn clone(&self) -> Self { + self.new_cloned() + } +} diff --git a/crates/evm/evm/src/executors/trace.rs b/crates/evm/evm/src/executors/trace.rs index ee93d0920..be0be748e 100644 --- a/crates/evm/evm/src/executors/trace.rs +++ b/crates/evm/evm/src/executors/trace.rs @@ -6,7 +6,7 @@ use foundry_evm_traces::{InternalTraceMode, TraceMode}; use revm::primitives::{Env, SpecId}; use std::ops::{Deref, DerefMut}; -use super::strategy::ExecutorStrategy; +use super::strategy::Strategy; /// A default executor with tracing enabled pub struct TracingExecutor { @@ -21,9 +21,9 @@ impl TracingExecutor { debug: bool, decode_internal: bool, alphanet: bool, - strategy: Box, + strategy: Strategy, ) -> Self { - let db = Backend::spawn(fork, strategy.new_backend_strategy()); + let db = Backend::spawn(fork, strategy.inner.new_backend_strategy()); let trace_mode = TraceMode::Call.with_debug(debug).with_decode_internal(if decode_internal { InternalTraceMode::Full diff --git a/crates/forge/bin/cmd/test/mod.rs b/crates/forge/bin/cmd/test/mod.rs index 384504bef..96335bdc7 100644 --- a/crates/forge/bin/cmd/test/mod.rs +++ b/crates/forge/bin/cmd/test/mod.rs @@ -368,7 +368,10 @@ impl TestArgs { // Prepare the test builder. let config = Arc::new(config); - strategy.zksync_set_dual_compiled_contracts(dual_compiled_contracts.unwrap_or_default()); + strategy.inner.zksync_set_dual_compiled_contracts( + strategy.context.as_mut(), + dual_compiled_contracts.unwrap_or_default(), + ); let runner = MultiContractRunnerBuilder::new(config.clone()) .set_debug(should_debug) diff --git a/crates/forge/src/multi_runner.rs b/crates/forge/src/multi_runner.rs index f2f901590..c716318da 100644 --- a/crates/forge/src/multi_runner.rs +++ b/crates/forge/src/multi_runner.rs @@ -18,7 +18,7 @@ use foundry_config::Config; use foundry_evm::{ backend::Backend, decode::RevertDecoder, - executors::{strategy::ExecutorStrategy, ExecutorBuilder}, + executors::{strategy::Strategy, ExecutorBuilder}, fork::CreateFork, inspectors::CheatsConfig, opts::EvmOpts, @@ -85,7 +85,7 @@ pub struct MultiContractRunner { /// Library addresses used to link contracts. pub libraries: Libraries, /// Execution strategy. - pub strategy: Box, + pub strategy: Strategy, } impl MultiContractRunner { @@ -178,7 +178,7 @@ impl MultiContractRunner { trace!("running all tests"); // The DB backend that serves all the data. - let db = Backend::spawn(self.fork.take(), self.strategy.new_backend_strategy()); + let db = Backend::spawn(self.fork.take(), self.strategy.inner.new_backend_strategy()); let find_timer = Instant::now(); let contracts = self.matching_contracts(filter).collect::>(); @@ -250,7 +250,7 @@ impl MultiContractRunner { Some(self.known_contracts.clone()), Some(artifact_id.name.clone()), Some(artifact_id.version.clone()), - self.strategy.new_cheatcode_inspector_strategy(), + self.strategy.inner.new_cheatcode_inspector_strategy(self.strategy.context.as_ref()), ); let trace_mode = TraceMode::default() @@ -270,7 +270,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, self.strategy.new_cloned()); + .build(self.env.clone(), db, self.strategy.clone()); if !enabled!(tracing::Level::TRACE) { span_name = get_contract_name(&identifier); @@ -407,7 +407,7 @@ impl MultiContractRunnerBuilder { zk_output: Option, env: revm::primitives::Env, evm_opts: EvmOpts, - strategy: Box, + strategy: Strategy, ) -> Result { let mut known_contracts = ContractsByArtifact::default(); let output = output.with_stripped_file_prefixes(root); diff --git a/crates/forge/src/runner.rs b/crates/forge/src/runner.rs index 067440c83..5f29c02d4 100644 --- a/crates/forge/src/runner.rs +++ b/crates/forge/src/runner.rs @@ -151,11 +151,7 @@ impl ContractRunner<'_> { // nonce. if let Some(cheatcodes) = &mut self.executor.inspector.cheatcodes { debug!("test contract deployed"); - cheatcodes - .strategy - .as_mut() - .expect("failed acquiring strategy") - .base_contract_deployed(); + cheatcodes.strategy.inner.base_contract_deployed(cheatcodes.strategy.context.as_mut()); } // Optionally call the `setUp` function diff --git a/crates/forge/tests/it/test_helpers.rs b/crates/forge/tests/it/test_helpers.rs index cd80f6d52..beb7e4f88 100644 --- a/crates/forge/tests/it/test_helpers.rs +++ b/crates/forge/tests/it/test_helpers.rs @@ -3,7 +3,7 @@ use alloy_chains::NamedChain; use alloy_primitives::U256; use forge::{ - executors::strategy::EvmExecutorStrategy, revm::primitives::SpecId, MultiContractRunner, + executors::strategy::new_evm_strategy, revm::primitives::SpecId, MultiContractRunner, MultiContractRunnerBuilder, TestOptions, TestOptionsBuilder, }; use foundry_cli::utils; @@ -360,7 +360,9 @@ impl ForgeTestData { let sender = zk_config.sender; let mut strategy = utils::get_executor_strategy(&zk_config); - strategy.zksync_set_dual_compiled_contracts(dual_compiled_contracts); + strategy + .inner + .zksync_set_dual_compiled_contracts(strategy.context.as_mut(), dual_compiled_contracts); let mut builder = self.base_runner(); builder.config = Arc::new(zk_config); builder @@ -382,7 +384,7 @@ impl ForgeTestData { None, opts.local_evm_env(), opts, - Box::new(EvmExecutorStrategy::default()), + new_evm_strategy(), ) .unwrap() } @@ -399,14 +401,7 @@ impl ForgeTestData { self.base_runner() .with_fork(fork) - .build( - self.project.root(), - self.output.clone(), - None, - env, - opts, - Box::new(EvmExecutorStrategy::default()), - ) + .build(self.project.root(), self.output.clone(), None, env, opts, new_evm_strategy()) .unwrap() } } diff --git a/crates/script/src/lib.rs b/crates/script/src/lib.rs index c56058bc4..80d437744 100644 --- a/crates/script/src/lib.rs +++ b/crates/script/src/lib.rs @@ -598,7 +598,7 @@ impl ScriptConfig { Some(db) => db.clone(), None => { let fork = self.evm_opts.get_fork(&self.config, env.clone()); - let backend = Backend::spawn(fork, strategy.new_backend_strategy()); + let backend = Backend::spawn(fork, strategy.inner.new_backend_strategy()); self.backends.insert(fork_url.clone(), backend.clone()); backend } @@ -607,7 +607,7 @@ impl ScriptConfig { // It's only really `None`, when we don't pass any `--fork-url`. And if so, there is // no need to cache it, since there won't be any onchain simulation that we'd need // to cache the backend for. - Backend::spawn(None, strategy.new_backend_strategy()) + Backend::spawn(None, strategy.inner.new_backend_strategy()) }; // We need to enable tracing to decode contract names: local or external. @@ -624,10 +624,13 @@ impl ScriptConfig { if let Some((known_contracts, script_wallets, target, dual_compiled_contracts)) = cheats_data { - strategy.zksync_set_dual_compiled_contracts(dual_compiled_contracts); + strategy.inner.zksync_set_dual_compiled_contracts( + strategy.context.as_mut(), + dual_compiled_contracts, + ); if let Some(fork_url) = &self.evm_opts.fork_url { - strategy.zksync_set_fork_env(fork_url, &env)?; + strategy.inner.zksync_set_fork_env(strategy.context.as_mut(), fork_url, &env)?; } builder = builder.inspectors(|stack| { @@ -639,7 +642,9 @@ impl ScriptConfig { Some(known_contracts), Some(target.name), Some(target.version), - strategy.new_cheatcode_inspector_strategy(), + strategy + .inner + .new_cheatcode_inspector_strategy(strategy.context.as_ref()), ) .into(), ) diff --git a/crates/script/src/runner.rs b/crates/script/src/runner.rs index 20ad42cf4..cb87bffc8 100644 --- a/crates/script/src/runner.rs +++ b/crates/script/src/runner.rs @@ -174,11 +174,7 @@ impl ScriptRunner { // nonce. if let Some(cheatcodes) = &mut self.executor.inspector.cheatcodes { debug!("script deployed"); - cheatcodes - .strategy - .as_mut() - .expect("failed acquiring strategy") - .base_contract_deployed(); + cheatcodes.strategy.inner.base_contract_deployed(cheatcodes.strategy.context.as_mut()); } // Optionally call the `setUp` function diff --git a/crates/strategy/zksync/src/backend.rs b/crates/strategy/zksync/src/backend.rs index ee75740a7..9ab91be12 100644 --- a/crates/strategy/zksync/src/backend.rs +++ b/crates/strategy/zksync/src/backend.rs @@ -1,7 +1,7 @@ -use std::collections::hash_map::Entry; +use std::{any::Any, collections::hash_map::Entry}; use alloy_primitives::{map::HashMap, Address, U256}; -use foundry_evm::backend::strategy::BackendStrategyExt; +use foundry_evm::backend::strategy::{BackendStrategyContext, BackendStrategyExt}; use foundry_evm_core::backend::{ strategy::{ BackendStrategy, BackendStrategyForkInfo, EvmBackendMergeStrategy, EvmBackendStrategy, @@ -18,12 +18,30 @@ use serde::{Deserialize, Serialize}; use tracing::trace; #[derive(Debug, Default, Clone, Serialize, Deserialize)] -pub struct ZksyncBackendStrategy { - evm: EvmBackendStrategy, +pub struct ZksyncBackendStrategyContext { /// Store storage keys per contract address for immutable variables. persistent_immutable_keys: HashMap>, } +impl BackendStrategyContext for ZksyncBackendStrategyContext { + fn new_cloned(&self) -> Box { + Box::new(self.clone()) + } + + fn as_any_ref(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + +#[derive(Debug, Default, Clone, Serialize, Deserialize)] +pub struct ZksyncBackendStrategy { + evm: EvmBackendStrategy, +} + #[derive(Debug, Default, Clone, Serialize, Deserialize)] pub struct ZkBackendInspectData { #[serde(skip_serializing_if = "Option::is_none")] @@ -35,6 +53,10 @@ pub struct ZkBackendInspectData { pub use_evm: bool, } +fn get_context(ctx: &mut dyn BackendStrategyContext) -> &mut ZksyncBackendStrategyContext { + ctx.as_any_mut().downcast_mut().expect("expected ZksyncBackendStrategyContext") +} + impl BackendStrategy for ZksyncBackendStrategy { fn name(&self) -> &'static str { "zk" @@ -47,13 +69,16 @@ impl BackendStrategy for ZksyncBackendStrategy { /// When creating or switching forks, we update the AccountInfo of the contract. fn update_fork_db( &self, + ctx: &mut dyn BackendStrategyContext, fork_info: BackendStrategyForkInfo<'_>, mem_db: &FoundryEvmInMemoryDB, backend_inner: &BackendInner, active_journaled_state: &mut JournaledState, target_fork: &mut Fork, ) { + let ctx = get_context(ctx); self.update_fork_db_contracts( + ctx, fork_info, mem_db, backend_inner, @@ -64,13 +89,20 @@ impl BackendStrategy for ZksyncBackendStrategy { fn merge_journaled_state_data( &self, + ctx: &mut dyn BackendStrategyContext, addr: Address, active_journaled_state: &JournaledState, fork_journaled_state: &mut JournaledState, ) { - self.evm.merge_journaled_state_data(addr, active_journaled_state, fork_journaled_state); + self.evm.merge_journaled_state_data( + ctx, + addr, + active_journaled_state, + fork_journaled_state, + ); + let ctx = get_context(ctx); let zk_state = - &ZksyncMergeState { persistent_immutable_keys: &self.persistent_immutable_keys }; + &ZksyncMergeState { persistent_immutable_keys: &ctx.persistent_immutable_keys }; ZksyncBackendMerge::merge_zk_journaled_state_data( addr, active_journaled_state, @@ -79,10 +111,17 @@ impl BackendStrategy for ZksyncBackendStrategy { ); } - fn merge_db_account_data(&self, addr: Address, active: &ForkDB, fork_db: &mut ForkDB) { - self.evm.merge_db_account_data(addr, active, fork_db); + fn merge_db_account_data( + &self, + ctx: &mut dyn BackendStrategyContext, + addr: Address, + active: &ForkDB, + fork_db: &mut ForkDB, + ) { + self.evm.merge_db_account_data(ctx, addr, active, fork_db); + let ctx = get_context(ctx); let zk_state = - &ZksyncMergeState { persistent_immutable_keys: &self.persistent_immutable_keys }; + &ZksyncMergeState { persistent_immutable_keys: &ctx.persistent_immutable_keys }; ZksyncBackendMerge::merge_zk_account_data(addr, active, fork_db, zk_state); } } @@ -91,6 +130,7 @@ impl ZksyncBackendStrategy { /// Merges the state of all `accounts` from the currently active db into the given `fork` pub(crate) fn update_fork_db_contracts( &self, + ctx: &mut ZksyncBackendStrategyContext, fork_info: BackendStrategyForkInfo<'_>, mem_db: &FoundryEvmInMemoryDB, backend_inner: &BackendInner, @@ -107,7 +147,7 @@ impl ZksyncBackendStrategy { let accounts = backend_inner.persistent_accounts.iter().copied(); let zk_state = - &ZksyncMergeState { persistent_immutable_keys: &self.persistent_immutable_keys }; + &ZksyncMergeState { persistent_immutable_keys: &ctx.persistent_immutable_keys }; if let Some(db) = fork_info.active_fork.map(|f| &f.db) { ZksyncBackendMerge::merge_account_data( accounts, @@ -129,8 +169,14 @@ impl ZksyncBackendStrategy { } impl BackendStrategyExt for ZksyncBackendStrategy { - fn zksync_save_immutable_storage(&mut self, addr: Address, keys: HashSet) { - self.persistent_immutable_keys + fn zksync_save_immutable_storage( + &self, + ctx: &mut dyn BackendStrategyContext, + addr: Address, + keys: HashSet, + ) { + let ctx = get_context(ctx); + ctx.persistent_immutable_keys .entry(addr) .and_modify(|entry| entry.extend(&keys)) .or_insert(keys); diff --git a/crates/strategy/zksync/src/cheatcode.rs b/crates/strategy/zksync/src/cheatcode.rs index b29f8a741..f8efc80c5 100644 --- a/crates/strategy/zksync/src/cheatcode.rs +++ b/crates/strategy/zksync/src/cheatcode.rs @@ -10,7 +10,8 @@ use alloy_sol_types::SolValue; use foundry_cheatcodes::{ journaled_account, make_acc_non_empty, strategy::{ - CheatcodeInspectorStrategy, CheatcodeInspectorStrategyExt, EvmCheatcodeInspectorStrategy, + CheatcodeInspectorStrategy, CheatcodeInspectorStrategyContext, + CheatcodeInspectorStrategyExt, EvmCheatcodeInspectorStrategy, }, Broadcast, BroadcastableTransaction, BroadcastableTransactions, Cheatcodes, CheatcodesExecutor, CheatsConfig, CheatsCtxt, CommonCreateInput, DealRecord, Ecx, Error, InnerEcx, Result, Vm, @@ -83,7 +84,10 @@ macro_rules! bail { #[derive(Debug, Default, Clone)] pub struct ZksyncCheatcodeInspectorStrategy { evm: EvmCheatcodeInspectorStrategy, +} +#[derive(Debug, Default, Clone)] +pub struct ZksyncCheatcodeInspectorStrategyContext { pub using_zk_vm: bool, /// When in zkEVM context, execute the next CALL or CREATE in the EVM instead. @@ -128,7 +132,7 @@ pub struct ZksyncCheatcodeInspectorStrategy { pub zk_env: ZkEnv, } -impl ZksyncCheatcodeInspectorStrategy { +impl ZksyncCheatcodeInspectorStrategyContext { pub fn new(dual_compiled_contracts: DualCompiledContracts, zk_env: ZkEnv) -> Self { // We add the empty bytecode manually so it is correctly translated in zk mode. // This is used in many places in foundry, e.g. in cheatcode contract's account code. @@ -168,7 +172,6 @@ impl ZksyncCheatcodeInspectorStrategy { persisted_factory_deps.insert(zk_bytecode_hash, zk_deployed_bytecode); Self { - evm: EvmCheatcodeInspectorStrategy::default(), using_zk_vm: false, // We need to migrate once on initialize_interp skip_zk_vm: false, skip_zk_vm_addresses: Default::default(), @@ -185,6 +188,32 @@ impl ZksyncCheatcodeInspectorStrategy { } } +impl CheatcodeInspectorStrategyContext for ZksyncCheatcodeInspectorStrategyContext { + fn new_cloned(&self) -> Box { + Box::new(self.clone()) + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn as_any_ref(&self) -> &dyn std::any::Any { + self + } +} + +fn get_context( + ctx: &mut dyn CheatcodeInspectorStrategyContext, +) -> &mut ZksyncCheatcodeInspectorStrategyContext { + ctx.as_any_mut().downcast_mut().expect("expected ZksyncCheatcodeInspectorStrategyContext") +} + +fn get_context_ref( + ctx: &dyn CheatcodeInspectorStrategyContext, +) -> &ZksyncCheatcodeInspectorStrategyContext { + ctx.as_any_ref().downcast_ref().expect("expected ZksyncCheatcodeInspectorStrategyContext") +} + /// Allows overriding nonce update behavior for the tx caller in the zkEVM. /// /// Since each CREATE or CALL is executed as a separate transaction within zkEVM, we currently skip @@ -227,8 +256,10 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { Box::new(self.clone()) } - fn get_nonce(&mut self, ccx: &mut CheatsCtxt<'_, '_, '_, '_>, address: Address) -> Result { - if !self.using_zk_vm { + fn get_nonce(&self, ccx: &mut CheatsCtxt<'_, '_, '_, '_>, address: Address) -> Result { + let ctx = get_context(ccx.state.strategy.context.as_mut()); + + if !ctx.using_zk_vm { return self.evm.get_nonce(ccx, address); } @@ -236,19 +267,23 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { Ok(nonce) } - fn base_contract_deployed(&mut self) { + fn base_contract_deployed(&self, ctx: &mut dyn CheatcodeInspectorStrategyContext) { + let ctx = get_context(ctx); + debug!("allowing startup storage migration"); - self.zk_startup_migration.allow(); + ctx.zk_startup_migration.allow(); debug!("allowing persisting next nonce update"); - self.zk_persist_nonce_update.persist_next(); + ctx.zk_persist_nonce_update.persist_next(); } fn cheatcode_get_nonce( - &mut self, + &self, ccx: &mut CheatsCtxt<'_, '_, '_, '_>, address: Address, ) -> foundry_cheatcodes::Result { - if !self.using_zk_vm { + let ctx = get_context(ccx.state.strategy.context.as_mut()); + + if !ctx.using_zk_vm { let nonce = self.evm.get_nonce(ccx, address)?; return Ok(nonce.abi_encode()); } @@ -257,8 +292,10 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { Ok(nonce.abi_encode()) } - fn cheatcode_roll(&mut self, ccx: &mut CheatsCtxt<'_, '_, '_, '_>, new_height: U256) -> Result { - if !self.using_zk_vm { + fn cheatcode_roll(&self, ccx: &mut CheatsCtxt<'_, '_, '_, '_>, new_height: U256) -> Result { + let ctx = get_context(ccx.state.strategy.context.as_mut()); + + if !ctx.using_zk_vm { return self.evm.cheatcode_roll(ccx, new_height); } @@ -267,12 +304,10 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { Ok(Default::default()) } - fn cheatcode_warp( - &mut self, - ccx: &mut CheatsCtxt<'_, '_, '_, '_>, - new_timestamp: U256, - ) -> Result { - if !self.using_zk_vm { + fn cheatcode_warp(&self, ccx: &mut CheatsCtxt<'_, '_, '_, '_>, new_timestamp: U256) -> Result { + let ctx = get_context(ccx.state.strategy.context.as_mut()); + + if !ctx.using_zk_vm { return self.evm.cheatcode_warp(ccx, new_timestamp); } @@ -282,12 +317,14 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { } fn cheatcode_deal( - &mut self, + &self, ccx: &mut CheatsCtxt<'_, '_, '_, '_>, address: Address, new_balance: U256, ) -> Result { - if !self.using_zk_vm { + let ctx = get_context(ccx.state.strategy.context.as_mut()); + + if !ctx.using_zk_vm { return self.evm.cheatcode_deal(ccx, address, new_balance); } @@ -298,12 +335,14 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { } fn cheatcode_etch( - &mut self, + &self, ccx: &mut CheatsCtxt<'_, '_, '_, '_>, target: Address, new_runtime_bytecode: &Bytes, ) -> Result { - if !self.using_zk_vm { + let ctx = get_context(ccx.state.strategy.context.as_mut()); + + if !ctx.using_zk_vm { return self.evm.cheatcode_etch(ccx, target, new_runtime_bytecode); } @@ -312,11 +351,13 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { } fn cheatcode_reset_nonce( - &mut self, + &self, ccx: &mut CheatsCtxt<'_, '_, '_, '_>, account: Address, ) -> Result { - if !self.using_zk_vm { + let ctx = get_context(ccx.state.strategy.context.as_mut()); + + if !ctx.using_zk_vm { return self.evm.cheatcode_reset_nonce(ccx, account); } @@ -325,12 +366,14 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { } fn cheatcode_set_nonce( - &mut self, + &self, ccx: &mut CheatsCtxt<'_, '_, '_, '_>, account: Address, new_nonce: u64, ) -> Result { - if !self.using_zk_vm { + let ctx = get_context(ccx.state.strategy.context.as_mut()); + + if !ctx.using_zk_vm { return self.evm.cheatcode_set_nonce(ccx, account, new_nonce); } @@ -348,12 +391,14 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { } fn cheatcode_set_nonce_unsafe( - &mut self, + &self, ccx: &mut CheatsCtxt<'_, '_, '_, '_>, account: Address, new_nonce: u64, ) -> Result { - if !self.using_zk_vm { + let ctx = get_context(ccx.state.strategy.context.as_mut()); + + if !ctx.using_zk_vm { return self.evm.cheatcode_set_nonce_unsafe(ccx, account, new_nonce); } @@ -362,13 +407,15 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { } fn cheatcode_mock_call( - &mut self, + &self, ccx: &mut CheatsCtxt<'_, '_, '_, '_>, callee: Address, data: &Bytes, return_data: &Bytes, ) -> Result { - if !self.using_zk_vm { + let ctx = get_context(ccx.state.strategy.context.as_mut()); + + if !ctx.using_zk_vm { return self.evm.cheatcode_mock_call(ccx, callee, data, return_data); } @@ -386,13 +433,15 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { } fn cheatcode_mock_call_revert( - &mut self, + &self, ccx: &mut CheatsCtxt<'_, '_, '_, '_>, callee: Address, data: &Bytes, revert_data: &Bytes, ) -> Result { - if !self.using_zk_vm { + let ctx = get_context(ccx.state.strategy.context.as_mut()); + + if !ctx.using_zk_vm { return self.evm.cheatcode_mock_call_revert(ccx, callee, data, revert_data); } @@ -411,9 +460,11 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { } fn get_artifact_code(&self, state: &Cheatcodes, path: &str, deployed: bool) -> Result { + let ctx = get_context_ref(state.strategy.context.as_ref()); + Ok(get_artifact_code( - &self.dual_compiled_contracts, - self.using_zk_vm, + &ctx.dual_compiled_contracts, + ctx.using_zk_vm, &state.config, path, deployed, @@ -422,15 +473,18 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { } fn record_broadcastable_create_transactions( - &mut self, + &self, + ctx: &mut dyn CheatcodeInspectorStrategyContext, config: Arc, input: &dyn CommonCreateInput, ecx_inner: InnerEcx<'_, '_, '_>, broadcast: &Broadcast, broadcastable_transactions: &mut BroadcastableTransactions, ) { - if !self.using_zk_vm { + let ctx_zk = get_context(ctx); + if !ctx_zk.using_zk_vm { return self.evm.record_broadcastable_create_transactions( + ctx, config, input, ecx_inner, @@ -439,13 +493,15 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { ); } + let ctx = ctx_zk; + let is_fixed_gas_limit = foundry_cheatcodes::check_if_fixed_gas_limit(ecx_inner, input.gas_limit()); let init_code = input.init_code(); let to = Some(TxKind::Call(CONTRACT_DEPLOYER_ADDRESS.to_address())); let mut nonce = foundry_zksync_core::nonce(broadcast.new_origin, ecx_inner) as u64; - let find_contract = self + let find_contract = ctx .dual_compiled_contracts .find_bytecode(&init_code.0) .unwrap_or_else(|| panic!("failed finding contract for {init_code:?}")); @@ -453,7 +509,7 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { let constructor_args = find_contract.constructor_args(); let contract = find_contract.contract(); - let factory_deps = self.dual_compiled_contracts.fetch_all_factory_deps(contract); + let factory_deps = ctx.dual_compiled_contracts.fetch_all_factory_deps(contract); let create_input = foundry_zksync_core::encode_create_params( &input.scheme().unwrap_or(CreateScheme::Create), @@ -464,21 +520,20 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { let mut zk_tx_factory_deps = factory_deps; - let paymaster_params = - self.paymaster_params.clone().map(|paymaster_data| PaymasterParams { - paymaster: paymaster_data.address.to_h160(), - paymaster_input: paymaster_data.input.to_vec(), - }); + let paymaster_params = ctx.paymaster_params.clone().map(|paymaster_data| PaymasterParams { + paymaster: paymaster_data.address.to_h160(), + paymaster_input: paymaster_data.input.to_vec(), + }); let rpc = ecx_inner.db.active_fork_url(); - let injected_factory_deps = self + let injected_factory_deps = ctx .zk_use_factory_deps .iter() .map(|contract| { get_artifact_code( - &self.dual_compiled_contracts, - self.using_zk_vm, + &ctx.dual_compiled_contracts, + ctx.using_zk_vm, &config, contract, false, @@ -543,7 +598,8 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { } fn record_broadcastable_call_transactions( - &mut self, + &self, + ctx: &mut dyn CheatcodeInspectorStrategyContext, config: Arc, call: &CallInputs, ecx_inner: InnerEcx<'_, '_, '_>, @@ -551,8 +607,11 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { broadcastable_transactions: &mut BroadcastableTransactions, active_delegation: &mut Option, ) { - if !self.using_zk_vm { + let ctx_zk = get_context(ctx); + + if !ctx_zk.using_zk_vm { return self.evm.record_broadcastable_call_transactions( + ctx, config, call, ecx_inner, @@ -562,19 +621,21 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { ); } + let ctx = ctx_zk; + let is_fixed_gas_limit = foundry_cheatcodes::check_if_fixed_gas_limit(ecx_inner, call.gas_limit); let nonce = foundry_zksync_core::nonce(broadcast.new_origin, ecx_inner) as u64; - let factory_deps = &mut self.set_deployer_call_input_factory_deps; - let injected_factory_deps = self + let factory_deps = &mut ctx.set_deployer_call_input_factory_deps; + let injected_factory_deps = ctx .zk_use_factory_deps .iter() .flat_map(|contract| { let artifact_code = get_artifact_code( - &self.dual_compiled_contracts, - self.using_zk_vm, + &ctx.dual_compiled_contracts, + ctx.using_zk_vm, &config, contract, false, @@ -584,17 +645,16 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { panic!("failed to get bytecode for factory deps contract {contract}") }) .to_vec(); - let res = self.dual_compiled_contracts.find_bytecode(&artifact_code).unwrap(); - self.dual_compiled_contracts.fetch_all_factory_deps(res.contract()) + let res = ctx.dual_compiled_contracts.find_bytecode(&artifact_code).unwrap(); + ctx.dual_compiled_contracts.fetch_all_factory_deps(res.contract()) }) .collect_vec(); factory_deps.extend(injected_factory_deps.clone()); - let paymaster_params = - self.paymaster_params.clone().map(|paymaster_data| PaymasterParams { - paymaster: paymaster_data.address.to_h160(), - paymaster_input: paymaster_data.input.to_vec(), - }); + let paymaster_params = ctx.paymaster_params.clone().map(|paymaster_data| PaymasterParams { + paymaster: paymaster_data.address.to_h160(), + paymaster_input: paymaster_data.input.to_vec(), + }); let factory_deps = if call.target_address == DEFAULT_CREATE2_DEPLOYER_ZKSYNC { // We shouldn't need factory_deps for CALLs factory_deps.clone() @@ -634,18 +694,32 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { debug!(target: "cheatcodes", tx=?broadcastable_transactions.back().unwrap(), "broadcastable call"); } - fn post_initialize_interp(&mut self, _interpreter: &mut Interpreter, ecx: Ecx<'_, '_, '_>) { - if self.zk_startup_migration.is_allowed() && !self.using_zk_vm { - self.select_zk_vm(ecx, None); - self.zk_startup_migration.done(); + fn post_initialize_interp( + &self, + ctx: &mut dyn CheatcodeInspectorStrategyContext, + _interpreter: &mut Interpreter, + ecx: Ecx<'_, '_, '_>, + ) { + let ctx = get_context(ctx); + + if ctx.zk_startup_migration.is_allowed() && !ctx.using_zk_vm { + self.select_zk_vm(ctx, ecx, None); + ctx.zk_startup_migration.done(); debug!("startup zkEVM storage migration completed"); } } /// Returns true if handled. - fn pre_step_end(&mut self, interpreter: &mut Interpreter, ecx: Ecx<'_, '_, '_>) -> bool { + fn pre_step_end( + &self, + ctx: &mut dyn CheatcodeInspectorStrategyContext, + interpreter: &mut Interpreter, + ecx: Ecx<'_, '_, '_>, + ) -> bool { // override address(x).balance retrieval to make it consistent between EraVM and EVM - if !self.using_zk_vm { + let ctx = get_context(ctx); + + if !ctx.using_zk_vm { return false; } @@ -680,29 +754,44 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { } impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { - fn zksync_cheatcode_skip_zkvm(&mut self) -> Result { - self.skip_zk_vm = true; + fn zksync_cheatcode_skip_zkvm( + &self, + ctx: &mut dyn CheatcodeInspectorStrategyContext, + ) -> Result { + let ctx = get_context(ctx); + + ctx.skip_zk_vm = true; Ok(Default::default()) } fn zksync_cheatcode_set_paymaster( - &mut self, + &self, + ctx: &mut dyn CheatcodeInspectorStrategyContext, paymaster_address: Address, paymaster_input: &Bytes, ) -> Result { - self.paymaster_params = + let ctx = get_context(ctx); + + ctx.paymaster_params = Some(ZkPaymasterData { address: paymaster_address, input: paymaster_input.clone() }); Ok(Default::default()) } - fn zksync_cheatcode_use_factory_deps(&mut self, name: String) -> foundry_cheatcodes::Result { + fn zksync_cheatcode_use_factory_deps( + &self, + ctx: &mut dyn CheatcodeInspectorStrategyContext, + name: String, + ) -> foundry_cheatcodes::Result { + let ctx = get_context(ctx); + info!("Adding factory dependency: {:?}", name); - self.zk_use_factory_deps.push(name); + ctx.zk_use_factory_deps.push(name); Ok(Default::default()) } fn zksync_cheatcode_register_contract( - &mut self, + &self, + ctx: &mut dyn CheatcodeInspectorStrategyContext, name: String, zk_bytecode_hash: FixedBytes<32>, zk_deployed_bytecode: Vec, @@ -711,6 +800,8 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { evm_deployed_bytecode: Vec, evm_bytecode: Vec, ) -> Result { + let ctx = get_context(ctx); + let new_contract = DualCompiledContract { name, zk_bytecode_hash: H256(zk_bytecode_hash.0), @@ -721,7 +812,7 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { evm_bytecode, }; - if let Some(existing) = self.dual_compiled_contracts.iter().find(|contract| { + if let Some(existing) = ctx.dual_compiled_contracts.iter().find(|contract| { contract.evm_bytecode_hash == new_contract.evm_bytecode_hash && contract.zk_bytecode_hash == new_contract.zk_bytecode_hash }) { @@ -729,33 +820,51 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { return Ok(Default::default()) } - self.dual_compiled_contracts.push(new_contract); + ctx.dual_compiled_contracts.push(new_contract); Ok(Default::default()) } - fn zksync_record_create_address(&mut self, outcome: &CreateOutcome) { - if self.record_next_create_address { - self.record_next_create_address = false; + fn zksync_record_create_address( + &self, + ctx: &mut dyn CheatcodeInspectorStrategyContext, + outcome: &CreateOutcome, + ) { + let ctx = get_context(ctx); + + if ctx.record_next_create_address { + ctx.record_next_create_address = false; if let Some(address) = outcome.address { - self.skip_zk_vm_addresses.insert(address); + ctx.skip_zk_vm_addresses.insert(address); } } } - fn zksync_sync_nonce(&mut self, sender: Address, nonce: u64, ecx: Ecx<'_, '_, '_>) { + fn zksync_sync_nonce( + &self, + _ctx: &mut dyn CheatcodeInspectorStrategyContext, + sender: Address, + nonce: u64, + ecx: Ecx<'_, '_, '_>, + ) { // NOTE(zk): We sync with the nonce changes to ensure that the nonce matches foundry_zksync_core::cheatcodes::set_nonce(sender, U256::from(nonce), ecx); } - fn zksync_set_deployer_call_input(&mut self, call: &mut CallInputs) { - self.set_deployer_call_input_factory_deps.clear(); - if call.target_address == DEFAULT_CREATE2_DEPLOYER && self.using_zk_vm { + fn zksync_set_deployer_call_input( + &self, + ctx: &mut dyn CheatcodeInspectorStrategyContext, + call: &mut CallInputs, + ) { + let ctx = get_context(ctx); + + ctx.set_deployer_call_input_factory_deps.clear(); + if call.target_address == DEFAULT_CREATE2_DEPLOYER && ctx.using_zk_vm { call.target_address = DEFAULT_CREATE2_DEPLOYER_ZKSYNC; call.bytecode_address = DEFAULT_CREATE2_DEPLOYER_ZKSYNC; let (salt, init_code) = call.input.split_at(32); - let find_contract = self + let find_contract = ctx .dual_compiled_contracts .find_bytecode(init_code) .unwrap_or_else(|| panic!("failed finding contract for {init_code:?}")); @@ -764,8 +873,8 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { let contract = find_contract.contract(); // store these for broadcast reasons - self.set_deployer_call_input_factory_deps = - self.dual_compiled_contracts.fetch_all_factory_deps(contract); + ctx.set_deployer_call_input_factory_deps = + ctx.dual_compiled_contracts.fetch_all_factory_deps(contract); let create_input = foundry_zksync_core::encode_create_params( &CreateScheme::Create2 { salt: U256::from_be_slice(salt) }, @@ -781,19 +890,21 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { /// If `Some` is returned then the result must be returned immediately, else the call must be /// handled in EVM. fn zksync_try_create( - &mut self, + &self, state: &mut Cheatcodes, ecx: Ecx<'_, '_, '_>, input: &dyn CommonCreateInput, executor: &mut dyn CheatcodesExecutor, ) -> Option { - if !self.using_zk_vm { + let ctx = get_context(state.strategy.context.as_mut()); + + if !ctx.using_zk_vm { return None; } - if self.skip_zk_vm { - self.skip_zk_vm = false; // handled the skip, reset flag - self.record_next_create_address = true; + if ctx.skip_zk_vm { + ctx.skip_zk_vm = false; // handled the skip, reset flag + ctx.record_next_create_address = true; info!("running create in EVM, instead of zkEVM (skipped)"); return None } @@ -822,7 +933,7 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { info!("running create in zkEVM"); - let find_contract = self + let find_contract = ctx .dual_compiled_contracts .find_bytecode(&init_code.0) .unwrap_or_else(|| panic!("failed finding contract for {init_code:?}")); @@ -836,14 +947,14 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { constructor_args.to_vec(), ); - let mut factory_deps = self.dual_compiled_contracts.fetch_all_factory_deps(contract); - let injected_factory_deps = self + let mut factory_deps = ctx.dual_compiled_contracts.fetch_all_factory_deps(contract); + let injected_factory_deps = ctx .zk_use_factory_deps .iter() .flat_map(|contract| { let artifact_code = get_artifact_code( - &self.dual_compiled_contracts, - self.using_zk_vm, + &ctx.dual_compiled_contracts, + ctx.using_zk_vm, &state.config, contract, false, @@ -853,25 +964,25 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { panic!("failed to get bytecode for injected factory deps contract {contract}") }) .to_vec(); - let res = self.dual_compiled_contracts.find_bytecode(&artifact_code).unwrap(); - self.dual_compiled_contracts.fetch_all_factory_deps(res.contract()) + let res = ctx.dual_compiled_contracts.find_bytecode(&artifact_code).unwrap(); + ctx.dual_compiled_contracts.fetch_all_factory_deps(res.contract()) }) .collect_vec(); factory_deps.extend(injected_factory_deps); // NOTE(zk): Clear injected factory deps so that they are not sent on further transactions - self.zk_use_factory_deps.clear(); + ctx.zk_use_factory_deps.clear(); tracing::debug!(contract = contract.name, "using dual compiled contract"); - let zk_persist_nonce_update = self.zk_persist_nonce_update.check(); + let zk_persist_nonce_update = ctx.zk_persist_nonce_update.check(); let ccx = foundry_zksync_core::vm::CheatcodeTracerContext { mocked_calls: state.mocked_calls.clone(), expected_calls: Some(&mut state.expected_calls), accesses: state.accesses.as_mut(), - persisted_factory_deps: Some(&mut self.persisted_factory_deps), - paymaster_data: self.paymaster_params.take(), + persisted_factory_deps: Some(&mut ctx.persisted_factory_deps), + paymaster_data: ctx.paymaster_params.take(), persist_nonce_update: state.broadcast.is_some() || zk_persist_nonce_update, - zk_env: self.zk_env.clone(), + zk_env: ctx.zk_env.clone(), }; let zk_create = foundry_zksync_core::vm::ZkCreateInputs { @@ -933,7 +1044,12 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { .to_ru256() }) .collect::>(); - ecx.db.get_strategy().zksync_save_immutable_storage(addr, keys); + let strategy = ecx.db.get_strategy(); + strategy.inner.zksync_save_immutable_storage( + strategy.context.as_mut(), + addr, + keys, + ); } } @@ -1000,24 +1116,25 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { /// If `Some` is returned then the result must be returned immediately, else the call must be /// handled in EVM. fn zksync_try_call( - &mut self, + &self, state: &mut Cheatcodes, ecx: Ecx<'_, '_, '_>, call: &CallInputs, executor: &mut dyn CheatcodesExecutor, ) -> Option { + let ctx = get_context(state.strategy.context.as_mut()); + // We need to clear them out for the next call. - let factory_deps = std::mem::take(&mut self.set_deployer_call_input_factory_deps); + let factory_deps = std::mem::take(&mut ctx.set_deployer_call_input_factory_deps); - if !self.using_zk_vm { + if !ctx.using_zk_vm { return None; } // also skip if the target was created during a zkEVM skip - self.skip_zk_vm = - self.skip_zk_vm || self.skip_zk_vm_addresses.contains(&call.target_address); - if self.skip_zk_vm { - self.skip_zk_vm = false; // handled the skip, reset flag + ctx.skip_zk_vm = ctx.skip_zk_vm || ctx.skip_zk_vm_addresses.contains(&call.target_address); + if ctx.skip_zk_vm { + ctx.skip_zk_vm = false; // handled the skip, reset flag info!("running create in EVM, instead of zkEVM (skipped) {:#?}", call); return None; } @@ -1036,20 +1153,20 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { } info!("running call in zkEVM {:#?}", call); - let zk_persist_nonce_update = self.zk_persist_nonce_update.check(); + let zk_persist_nonce_update = ctx.zk_persist_nonce_update.check(); // NOTE(zk): Clear injected factory deps here even though it's actually used in broadcast. // To be consistent with where we clear factory deps in try_create_in_zk. - self.zk_use_factory_deps.clear(); + ctx.zk_use_factory_deps.clear(); let ccx = foundry_zksync_core::vm::CheatcodeTracerContext { mocked_calls: state.mocked_calls.clone(), expected_calls: Some(&mut state.expected_calls), accesses: state.accesses.as_mut(), - persisted_factory_deps: Some(&mut self.persisted_factory_deps), - paymaster_data: self.paymaster_params.take(), + persisted_factory_deps: Some(&mut ctx.persisted_factory_deps), + paymaster_data: ctx.paymaster_params.take(), persist_nonce_update: state.broadcast.is_some() || zk_persist_nonce_update, - zk_env: self.zk_env.clone(), + zk_env: ctx.zk_env.clone(), }; let mut gas = Gas::new(call.gas_limit); @@ -1155,15 +1272,27 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { } } - fn zksync_select_fork_vm(&mut self, data: InnerEcx<'_, '_, '_>, fork_id: LocalForkId) { - self.select_fork_vm(data, fork_id); + fn zksync_select_fork_vm( + &self, + ctx: &mut dyn CheatcodeInspectorStrategyContext, + data: InnerEcx<'_, '_, '_>, + fork_id: LocalForkId, + ) { + let ctx = get_context(ctx); + self.select_fork_vm(ctx, data, fork_id); } - fn zksync_cheatcode_select_zk_vm(&mut self, data: InnerEcx<'_, '_, '_>, enable: bool) { + fn zksync_cheatcode_select_zk_vm( + &self, + ctx: &mut dyn CheatcodeInspectorStrategyContext, + data: InnerEcx<'_, '_, '_>, + enable: bool, + ) { + let ctx = get_context(ctx); if enable { - self.select_zk_vm(data, None) + self.select_zk_vm(ctx, data, None) } else { - self.select_evm(data); + self.select_evm(ctx, data); } } } @@ -1175,25 +1304,34 @@ impl ZksyncCheatcodeInspectorStrategy { /// Additionally: /// * Translates block information /// * Translates all persisted addresses - pub fn select_fork_vm(&mut self, data: InnerEcx<'_, '_, '_>, fork_id: LocalForkId) { + pub fn select_fork_vm( + &self, + ctx: &mut ZksyncCheatcodeInspectorStrategyContext, + data: InnerEcx<'_, '_, '_>, + fork_id: LocalForkId, + ) { let fork_info = data.db.get_fork_info(fork_id).expect("failed getting fork info"); if fork_info.fork_type.is_evm() { - self.select_evm(data) + self.select_evm(ctx, data) } else { - self.select_zk_vm(data, Some(&fork_info.fork_env)) + self.select_zk_vm(ctx, data, Some(&fork_info.fork_env)) } } /// Switch to EVM and translate block info, balances, nonces and deployed codes for persistent /// accounts - pub fn select_evm(&mut self, data: InnerEcx<'_, '_, '_>) { - if !self.using_zk_vm { + pub fn select_evm( + &self, + ctx: &mut ZksyncCheatcodeInspectorStrategyContext, + data: InnerEcx<'_, '_, '_>, + ) { + if !ctx.using_zk_vm { tracing::info!("already in EVM"); return } tracing::info!("switching to EVM"); - self.using_zk_vm = false; + ctx.using_zk_vm = false; let system_account = SYSTEM_CONTEXT_ADDRESS.to_address(); journaled_account(data, system_account).expect("failed to load account"); @@ -1230,7 +1368,7 @@ impl ZksyncCheatcodeInspectorStrategy { .sload(account_code_account, account_code_key) .ok() .and_then(|zk_bytecode_hash| { - self.dual_compiled_contracts + ctx.dual_compiled_contracts .find_by_zk_bytecode_hash(zk_bytecode_hash.to_h256()) .map(|contract| { ( @@ -1258,14 +1396,19 @@ impl ZksyncCheatcodeInspectorStrategy { /// Switch to ZK-VM and translate block info, balances, nonces and deployed codes for persistent /// accounts - pub fn select_zk_vm(&mut self, data: InnerEcx<'_, '_, '_>, new_env: Option<&Env>) { - if self.using_zk_vm { + pub fn select_zk_vm( + &self, + ctx: &mut ZksyncCheatcodeInspectorStrategyContext, + data: InnerEcx<'_, '_, '_>, + new_env: Option<&Env>, + ) { + if ctx.using_zk_vm { tracing::info!("already in ZK-VM"); return } tracing::info!("switching to ZK-VM"); - self.using_zk_vm = true; + ctx.using_zk_vm = true; let env = new_env.unwrap_or(data.env.as_ref()); @@ -1303,7 +1446,7 @@ impl ZksyncCheatcodeInspectorStrategy { continue; } - if let Some(contract) = self.dual_compiled_contracts.iter().find(|contract| { + if let Some(contract) = ctx.dual_compiled_contracts.iter().find(|contract| { info.code_hash != KECCAK_EMPTY && info.code_hash == contract.evm_bytecode_hash }) { account_code_storage.insert( diff --git a/crates/strategy/zksync/src/executor.rs b/crates/strategy/zksync/src/executor.rs index e96c4ff60..a469d2f05 100644 --- a/crates/strategy/zksync/src/executor.rs +++ b/crates/strategy/zksync/src/executor.rs @@ -2,12 +2,13 @@ use alloy_primitives::{Address, U256}; use alloy_rpc_types::serde_helpers::OtherFields; use alloy_zksync::provider::{zksync_provider, ZksyncProvider}; use eyre::Result; -use foundry_cheatcodes::strategy::CheatcodeInspectorStrategy; use foundry_evm::{ - backend::{strategy::BackendStrategy, BackendResult, DatabaseExt}, + backend::{BackendResult, DatabaseExt}, executors::{ - strategy::{EvmExecutorStrategy, ExecutorStrategy, ExecutorStrategyExt}, + strategy::{ + EvmExecutorStrategy, ExecutorStrategy, ExecutorStrategyContext, ExecutorStrategyExt, + }, Executor, }, InspectorExt, @@ -21,19 +22,46 @@ use revm::{ use zksync_types::H256; use crate::{ - cheatcode::ZKSYNC_TRANSACTION_OTHER_FIELDS_KEY, ZksyncBackendStrategy, - ZksyncCheatcodeInspectorStrategy, + backend::ZksyncBackendStrategyContext, + cheatcode::{ZksyncCheatcodeInspectorStrategyContext, ZKSYNC_TRANSACTION_OTHER_FIELDS_KEY}, + ZksyncBackendStrategy, ZksyncCheatcodeInspectorStrategy, }; #[derive(Debug, Default, Clone)] -pub struct ZksyncExecutorStrategy { - evm: EvmExecutorStrategy, +pub struct ZksyncExecutorStrategyContext { inspect_context: Option, persisted_factory_deps: HashMap>, dual_compiled_contracts: DualCompiledContracts, zk_env: ZkEnv, } +impl ExecutorStrategyContext for ZksyncExecutorStrategyContext { + fn new_cloned(&self) -> Box { + Box::new(self.clone()) + } + + fn as_any_ref(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } +} + +#[derive(Debug, Default, Clone)] +pub struct ZksyncExecutorStrategy { + evm: EvmExecutorStrategy, +} + +fn get_context_ref(ctx: &dyn ExecutorStrategyContext) -> &ZksyncExecutorStrategyContext { + ctx.as_any_ref().downcast_ref().expect("expected ZksyncExecutorStrategyContext") +} + +fn get_context(ctx: &mut dyn ExecutorStrategyContext) -> &mut ZksyncExecutorStrategyContext { + ctx.as_any_mut().downcast_mut().expect("expected ZksyncExecutorStrategyContext") +} + impl ExecutorStrategy for ZksyncExecutorStrategy { fn name(&self) -> &'static str { "zk" @@ -43,13 +71,18 @@ impl ExecutorStrategy for ZksyncExecutorStrategy { Box::new(self.clone()) } - fn set_inspect_context(&mut self, other_fields: OtherFields) { + fn set_inspect_context( + &self, + ctx: &mut dyn ExecutorStrategyContext, + other_fields: OtherFields, + ) { + let ctx = get_context(ctx); let maybe_context = get_zksync_transaction_metadata(&other_fields); - self.inspect_context = maybe_context; + ctx.inspect_context = maybe_context; } fn set_balance( - &mut self, + &self, executor: &mut Executor, address: Address, amount: U256, @@ -63,7 +96,7 @@ impl ExecutorStrategy for ZksyncExecutorStrategy { } fn set_nonce( - &mut self, + &self, executor: &mut Executor, address: Address, nonce: u64, @@ -81,45 +114,60 @@ impl ExecutorStrategy for ZksyncExecutorStrategy { Ok(()) } - fn new_backend_strategy(&self) -> Box { - Box::new(ZksyncBackendStrategy::default()) + fn new_backend_strategy(&self) -> foundry_evm_core::backend::strategy::Strategy { + foundry_evm_core::backend::strategy::Strategy { + inner: Box::new(ZksyncBackendStrategy::default()), + context: Box::new(ZksyncBackendStrategyContext::default()), + } } - fn new_cheatcode_inspector_strategy(&self) -> Box { - Box::new(ZksyncCheatcodeInspectorStrategy::new( - self.dual_compiled_contracts.clone(), - self.zk_env.clone(), - )) + fn new_cheatcode_inspector_strategy( + &self, + ctx: &dyn ExecutorStrategyContext, + ) -> foundry_cheatcodes::strategy::Strategy { + let ctx = get_context_ref(ctx); + foundry_cheatcodes::strategy::Strategy { + inner: Box::new(ZksyncCheatcodeInspectorStrategy::default()), + context: Box::new(ZksyncCheatcodeInspectorStrategyContext::new( + ctx.dual_compiled_contracts.clone(), + ctx.zk_env.clone(), + )), + } } fn call_inspect( &self, + ctx: &dyn ExecutorStrategyContext, db: &mut dyn DatabaseExt, env: &mut EnvWithHandlerCfg, inspector: &mut dyn InspectorExt, ) -> eyre::Result { - match self.inspect_context.as_ref() { - None => self.evm.call_inspect(db, env, inspector), + let ctx_zk = get_context_ref(ctx); + match ctx_zk.inspect_context.as_ref() { + None => self.evm.call_inspect(ctx, db, env, inspector), Some(zk_tx) => foundry_zksync_core::vm::transact( - Some(&mut self.persisted_factory_deps.clone()), + Some(&mut ctx_zk.persisted_factory_deps.clone()), Some(zk_tx.factory_deps.clone()), zk_tx.paymaster_data.clone(), env, - &self.zk_env, + &ctx_zk.zk_env, db, ), } } fn transact_inspect( - &mut self, + &self, + ctx: &mut dyn ExecutorStrategyContext, db: &mut dyn DatabaseExt, env: &mut EnvWithHandlerCfg, executor_env: &EnvWithHandlerCfg, inspector: &mut dyn InspectorExt, ) -> eyre::Result { - match self.inspect_context.take() { - None => self.evm.transact_inspect(db, env, executor_env, inspector), + let ctx_zk = get_context(ctx); + + match ctx_zk.inspect_context.take() { + None => self.evm.transact_inspect(ctx, db, env, executor_env, inspector), Some(zk_tx) => { // apply fork-related env instead of cheatcode handler // since it won't be set by zkEVM @@ -127,11 +175,11 @@ impl ExecutorStrategy for ZksyncExecutorStrategy { env.tx.gas_price = executor_env.tx.gas_price; foundry_zksync_core::vm::transact( - Some(&mut self.persisted_factory_deps), + Some(&mut ctx_zk.persisted_factory_deps), Some(zk_tx.factory_deps), zk_tx.paymaster_data, env, - &self.zk_env, + &ctx_zk.zk_env, db, ) } @@ -141,13 +189,23 @@ impl ExecutorStrategy for ZksyncExecutorStrategy { impl ExecutorStrategyExt for ZksyncExecutorStrategy { fn zksync_set_dual_compiled_contracts( - &mut self, + &self, + ctx: &mut dyn ExecutorStrategyContext, dual_compiled_contracts: DualCompiledContracts, ) { - self.dual_compiled_contracts = dual_compiled_contracts; + let ctx = get_context(ctx); + ctx.dual_compiled_contracts = dual_compiled_contracts; + println!("SAVE {}", ctx.dual_compiled_contracts.len()); } - fn zksync_set_fork_env(&mut self, fork_url: &str, env: &Env) -> Result<()> { + fn zksync_set_fork_env( + &self, + ctx: &mut dyn ExecutorStrategyContext, + fork_url: &str, + env: &Env, + ) -> Result<()> { + let ctx = get_context(ctx); + let provider = zksync_provider().with_recommended_fillers().on_http(fork_url.parse()?); let block_number = env.block.number.try_into()?; // TODO(zk): switch to getFeeParams call when it is implemented for anvil-zksync @@ -158,7 +216,7 @@ impl ExecutorStrategyExt for ZksyncExecutorStrategy { .flatten(); if let Some(block_details) = maybe_block_details { - self.zk_env = ZkEnv { + ctx.zk_env = ZkEnv { l1_gas_price: block_details .l1_gas_price .try_into() @@ -192,3 +250,10 @@ pub fn get_zksync_transaction_metadata( .ok() .flatten() } + +pub fn new_zkysnc_strategy() -> foundry_evm::executors::strategy::Strategy { + foundry_evm::executors::strategy::Strategy { + inner: Box::new(ZksyncExecutorStrategy::default()), + context: Box::new(ZksyncExecutorStrategyContext::default()), + } +} diff --git a/crates/strategy/zksync/src/lib.rs b/crates/strategy/zksync/src/lib.rs index 0d7e2493b..8018e89be 100644 --- a/crates/strategy/zksync/src/lib.rs +++ b/crates/strategy/zksync/src/lib.rs @@ -11,4 +11,4 @@ mod executor; pub use backend::ZksyncBackendStrategy; pub use cheatcode::ZksyncCheatcodeInspectorStrategy; -pub use executor::{get_zksync_transaction_metadata, ZksyncExecutorStrategy}; +pub use executor::{get_zksync_transaction_metadata, new_zkysnc_strategy, ZksyncExecutorStrategy}; diff --git a/crates/verify/src/bytecode.rs b/crates/verify/src/bytecode.rs index 24ceaa547..5a9592b4b 100644 --- a/crates/verify/src/bytecode.rs +++ b/crates/verify/src/bytecode.rs @@ -240,7 +240,7 @@ impl VerifyBytecodeArgs { gen_blk_num, etherscan_metadata.evm_version()?.unwrap_or(EvmVersion::default()), evm_opts, - strategy.new_cloned(), + strategy.clone(), ) .await?; @@ -444,7 +444,7 @@ impl VerifyBytecodeArgs { simulation_block - 1, // env.fork_block_number etherscan_metadata.evm_version()?.unwrap_or(EvmVersion::default()), evm_opts, - strategy.new_cloned(), + strategy.clone(), ) .await?; env.block.number = U256::from(simulation_block); diff --git a/crates/verify/src/utils.rs b/crates/verify/src/utils.rs index fbf6a2664..60ba29715 100644 --- a/crates/verify/src/utils.rs +++ b/crates/verify/src/utils.rs @@ -14,7 +14,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::Strategy, TracingExecutor}, opts::EvmOpts, }; use reqwest::Url; @@ -325,7 +325,7 @@ pub async fn get_tracing_executor( fork_blk_num: u64, evm_version: EvmVersion, evm_opts: EvmOpts, - strategy: Box, + strategy: Strategy, ) -> Result<(Env, TracingExecutor)> { fork_config.fork_block_number = Some(fork_blk_num); fork_config.evm_version = evm_version; diff --git a/crates/zksync/compiler/src/zksolc/mod.rs b/crates/zksync/compiler/src/zksolc/mod.rs index f4e3b0699..ffe7349ed 100644 --- a/crates/zksync/compiler/src/zksolc/mod.rs +++ b/crates/zksync/compiler/src/zksolc/mod.rs @@ -299,4 +299,9 @@ impl DualCompiledContracts { pub fn push(&mut self, contract: DualCompiledContract) { self.contracts.push(contract); } + + /// Retrieves the length of the collection. + pub fn len(&self) -> usize { + self.contracts.len() + } } From 22e704c9651adc77c1352fa8e207390d45975bf6 Mon Sep 17 00:00:00 2001 From: Nisheeth Barthwal Date: Sat, 21 Dec 2024 22:40:16 +0100 Subject: [PATCH 2/6] clippy --- crates/zksync/compiler/src/zksolc/mod.rs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/crates/zksync/compiler/src/zksolc/mod.rs b/crates/zksync/compiler/src/zksolc/mod.rs index ffe7349ed..50eee0f2f 100644 --- a/crates/zksync/compiler/src/zksolc/mod.rs +++ b/crates/zksync/compiler/src/zksolc/mod.rs @@ -304,4 +304,9 @@ impl DualCompiledContracts { pub fn len(&self) -> usize { self.contracts.len() } + + /// Retrieves if the collection is empty. + pub fn is_empty(&self) -> bool { + self.contracts.is_empty() + } } From b83bf046eeb8755c908582bfc4332d707c961199 Mon Sep 17 00:00:00 2001 From: Nisheeth Barthwal Date: Sat, 21 Dec 2024 22:59:48 +0100 Subject: [PATCH 3/6] clippy --- crates/cheatcodes/src/strategy.rs | 3 ++- crates/evm/core/src/backend/strategy.rs | 2 +- crates/evm/evm/src/executors/strategy.rs | 2 +- crates/strategy/zksync/src/executor.rs | 1 - 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/cheatcodes/src/strategy.rs b/crates/cheatcodes/src/strategy.rs index 6424b2f44..dc1928b40 100644 --- a/crates/cheatcodes/src/strategy.rs +++ b/crates/cheatcodes/src/strategy.rs @@ -26,7 +26,7 @@ pub trait CheatcodeInspectorStrategyContext: Debug + Send + Sync + Any { impl CheatcodeInspectorStrategyContext for () { fn new_cloned(&self) -> Box { - Box::new(self.clone()) + Box::new(*self) } fn as_any_mut(&mut self) -> &mut dyn Any { @@ -195,6 +195,7 @@ pub trait CheatcodeInspectorStrategy: Debug + Send + Sync + CheatcodeInspectorSt ); /// Record broadcastable transaction during CALL. + #[allow(clippy::too_many_arguments)] fn record_broadcastable_call_transactions( &self, _ctx: &mut dyn CheatcodeInspectorStrategyContext, diff --git a/crates/evm/core/src/backend/strategy.rs b/crates/evm/core/src/backend/strategy.rs index 74dd3023a..b54eab781 100644 --- a/crates/evm/core/src/backend/strategy.rs +++ b/crates/evm/core/src/backend/strategy.rs @@ -19,7 +19,7 @@ pub trait BackendStrategyContext: Debug + Send + Sync + Any { impl BackendStrategyContext for () { fn new_cloned(&self) -> Box { - Box::new(self.clone()) + Box::new(()) } fn as_any_ref(&self) -> &dyn Any { diff --git a/crates/evm/evm/src/executors/strategy.rs b/crates/evm/evm/src/executors/strategy.rs index a2cc03cd2..2d5fcfd84 100644 --- a/crates/evm/evm/src/executors/strategy.rs +++ b/crates/evm/evm/src/executors/strategy.rs @@ -24,7 +24,7 @@ pub trait ExecutorStrategyContext: Debug + Send + Sync + Any { impl ExecutorStrategyContext for () { fn new_cloned(&self) -> Box { - Box::new(self.clone()) + Box::new(*self) } fn as_any_ref(&self) -> &dyn Any { diff --git a/crates/strategy/zksync/src/executor.rs b/crates/strategy/zksync/src/executor.rs index a469d2f05..824080792 100644 --- a/crates/strategy/zksync/src/executor.rs +++ b/crates/strategy/zksync/src/executor.rs @@ -195,7 +195,6 @@ impl ExecutorStrategyExt for ZksyncExecutorStrategy { ) { let ctx = get_context(ctx); ctx.dual_compiled_contracts = dual_compiled_contracts; - println!("SAVE {}", ctx.dual_compiled_contracts.len()); } fn zksync_set_fork_env( From 16d9cfab70dd58bf6b1efc9d13def77ba7b367f4 Mon Sep 17 00:00:00 2001 From: Nisheeth Barthwal Date: Sun, 22 Dec 2024 20:27:14 +0100 Subject: [PATCH 4/6] cleanup --- crates/cheatcodes/src/config.rs | 13 +-- crates/cheatcodes/src/evm.rs | 16 +-- crates/cheatcodes/src/evm/fork.rs | 4 +- crates/cheatcodes/src/evm/mock.rs | 4 +- crates/cheatcodes/src/fs.rs | 2 +- crates/cheatcodes/src/inspector.rs | 23 ++--- crates/cheatcodes/src/strategy.rs | 46 ++++++--- crates/cheatcodes/src/test.rs | 10 +- crates/chisel/src/executor.rs | 4 +- crates/cli/src/utils/mod.rs | 10 +- crates/evm/core/src/backend/cow.rs | 4 +- crates/evm/core/src/backend/mod.rs | 30 +++--- crates/evm/core/src/backend/strategy.rs | 42 +++++--- crates/evm/evm/src/executors/builder.rs | 4 +- crates/evm/evm/src/executors/mod.rs | 16 +-- crates/evm/evm/src/executors/strategy.rs | 57 ++++++----- crates/evm/evm/src/executors/trace.rs | 6 +- crates/forge/bin/cmd/test/mod.rs | 2 +- crates/forge/src/multi_runner.rs | 10 +- crates/forge/src/runner.rs | 2 +- crates/forge/tests/it/test_helpers.rs | 2 +- crates/script/src/lib.rs | 10 +- crates/script/src/runner.rs | 2 +- crates/strategy/zksync/src/backend.rs | 55 ++++++----- crates/strategy/zksync/src/cheatcode.rs | 121 +++++++++++++---------- crates/strategy/zksync/src/executor.rs | 58 ++++++----- crates/strategy/zksync/src/lib.rs | 8 +- crates/verify/src/utils.rs | 4 +- 28 files changed, 318 insertions(+), 247 deletions(-) diff --git a/crates/cheatcodes/src/config.rs b/crates/cheatcodes/src/config.rs index 9c519ee93..5a32f40c9 100644 --- a/crates/cheatcodes/src/config.rs +++ b/crates/cheatcodes/src/config.rs @@ -1,8 +1,5 @@ use super::Result; -use crate::{ - strategy::{new_evm_strategy, Strategy}, - Vm::Rpc, -}; +use crate::{strategy::CheatcodeInspectorStrategy, Vm::Rpc}; use alloy_primitives::{map::AddressHashMap, U256}; use foundry_common::{fs::normalize_path, ContractsByArtifact}; use foundry_compilers::{utils::canonicalize, ProjectPathsConfig}; @@ -57,7 +54,7 @@ pub struct CheatsConfig { /// Version of the script/test contract which is currently running. pub running_version: Option, /// The behavior strategy. - pub strategy: Strategy, + pub strategy: CheatcodeInspectorStrategy, /// Whether to enable legacy (non-reverting) assertions. pub assertions_revert: bool, /// Optional seed for the RNG algorithm. @@ -73,7 +70,7 @@ impl CheatsConfig { available_artifacts: Option, running_contract: Option, running_version: Option, - strategy: Strategy, + strategy: CheatcodeInspectorStrategy, ) -> Self { let mut allowed_paths = vec![config.root.0.clone()]; allowed_paths.extend(config.libs.clone()); @@ -234,7 +231,7 @@ impl Default for CheatsConfig { available_artifacts: Default::default(), running_contract: Default::default(), running_version: Default::default(), - strategy: new_evm_strategy(), + strategy: CheatcodeInspectorStrategy::new_evm(), assertions_revert: true, seed: None, } @@ -253,7 +250,7 @@ mod tests { None, None, None, - new_evm_strategy(), + CheatcodeInspectorStrategy::new_evm(), ) } diff --git a/crates/cheatcodes/src/evm.rs b/crates/cheatcodes/src/evm.rs index 43eb9f50b..e5d6d3db5 100644 --- a/crates/cheatcodes/src/evm.rs +++ b/crates/cheatcodes/src/evm.rs @@ -64,7 +64,7 @@ impl Cheatcode for getNonce_0Call { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { account } = self; - ccx.state.strategy.inner.clone().cheatcode_get_nonce(ccx, *account) + ccx.state.strategy.runner.clone().cheatcode_get_nonce(ccx, *account) } } @@ -350,7 +350,7 @@ impl Cheatcode for getBlobhashesCall { impl Cheatcode for rollCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { newHeight } = self; - ccx.state.strategy.inner.clone().cheatcode_roll(ccx, *newHeight) + ccx.state.strategy.runner.clone().cheatcode_roll(ccx, *newHeight) } } @@ -372,7 +372,7 @@ impl Cheatcode for txGasPriceCall { impl Cheatcode for warpCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { newTimestamp } = self; - ccx.state.strategy.inner.clone().cheatcode_warp(ccx, *newTimestamp) + ccx.state.strategy.runner.clone().cheatcode_warp(ccx, *newTimestamp) } } @@ -407,7 +407,7 @@ impl Cheatcode for dealCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { account: address, newBalance: new_balance } = *self; - ccx.state.strategy.inner.clone().cheatcode_deal(ccx, address, new_balance) + ccx.state.strategy.runner.clone().cheatcode_deal(ccx, address, new_balance) } } @@ -415,14 +415,14 @@ impl Cheatcode for etchCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { target, newRuntimeBytecode } = self; - ccx.state.strategy.inner.clone().cheatcode_etch(ccx, *target, newRuntimeBytecode) + ccx.state.strategy.runner.clone().cheatcode_etch(ccx, *target, newRuntimeBytecode) } } impl Cheatcode for resetNonceCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { account } = self; - ccx.state.strategy.inner.clone().cheatcode_reset_nonce(ccx, *account) + ccx.state.strategy.runner.clone().cheatcode_reset_nonce(ccx, *account) } } @@ -430,7 +430,7 @@ impl Cheatcode for setNonceCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { account, newNonce } = *self; - ccx.state.strategy.inner.clone().cheatcode_set_nonce(ccx, account, newNonce) + ccx.state.strategy.runner.clone().cheatcode_set_nonce(ccx, account, newNonce) } } @@ -438,7 +438,7 @@ impl Cheatcode for setNonceUnsafeCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { account, newNonce } = *self; - ccx.state.strategy.inner.clone().cheatcode_set_nonce_unsafe(ccx, account, newNonce) + ccx.state.strategy.runner.clone().cheatcode_set_nonce_unsafe(ccx, account, newNonce) } } diff --git a/crates/cheatcodes/src/evm/fork.rs b/crates/cheatcodes/src/evm/fork.rs index a1bcc9ccb..e1db17a3d 100644 --- a/crates/cheatcodes/src/evm/fork.rs +++ b/crates/cheatcodes/src/evm/fork.rs @@ -125,7 +125,7 @@ impl Cheatcode for selectForkCall { persist_caller(ccx); check_broadcast(ccx.state)?; - ccx.state.strategy.inner.zksync_select_fork_vm( + ccx.state.strategy.runner.zksync_select_fork_vm( ccx.state.strategy.context.as_mut(), ccx.ecx, *forkId, @@ -285,7 +285,7 @@ fn create_select_fork(ccx: &mut CheatsCtxt, url_or_alias: &str, block: Option Result { let Self { callee, data, returnData } = self; - ccx.state.strategy.inner.clone().cheatcode_mock_call(ccx, *callee, data, returnData) + ccx.state.strategy.runner.clone().cheatcode_mock_call(ccx, *callee, data, returnData) } } @@ -83,7 +83,7 @@ impl Cheatcode for mockCalls_1Call { impl Cheatcode for mockCallRevert_0Call { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { callee, data, revertData } = self; - ccx.state.strategy.inner.clone().cheatcode_mock_call_revert(ccx, *callee, data, revertData) + ccx.state.strategy.runner.clone().cheatcode_mock_call_revert(ccx, *callee, data, revertData) } } diff --git a/crates/cheatcodes/src/fs.rs b/crates/cheatcodes/src/fs.rs index b99ed2945..96e53704c 100644 --- a/crates/cheatcodes/src/fs.rs +++ b/crates/cheatcodes/src/fs.rs @@ -283,7 +283,7 @@ impl Cheatcode for getArtifactPathByDeployedCodeCall { impl Cheatcode for getCodeCall { fn apply(&self, state: &mut Cheatcodes) -> Result { let Self { artifactPath: path } = self; - state.strategy.inner.get_artifact_code(state, path, false) + state.strategy.runner.get_artifact_code(state, path, false) } } diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index 7ae10782c..ad1f757e2 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -7,7 +7,7 @@ use crate::{ DealRecord, GasRecord, }, script::{Broadcast, Wallets}, - strategy::Strategy, + strategy::CheatcodeInspectorStrategy, test::{ assume::AssumeNoRevert, expect::{self, ExpectedEmit, ExpectedRevert, ExpectedRevertKind}, @@ -527,7 +527,7 @@ pub struct Cheatcodes { pub wallets: Option, /// The behavior strategy. - pub strategy: Strategy, + pub strategy: CheatcodeInspectorStrategy, } impl Clone for Cheatcodes { @@ -762,7 +762,7 @@ impl Cheatcodes { if ecx_inner.journaled_state.depth() == broadcast.depth { input.set_caller(broadcast.new_origin); - self.strategy.inner.record_broadcastable_create_transactions( + self.strategy.runner.record_broadcastable_create_transactions( self.strategy.context.as_mut(), self.config.clone(), &input, @@ -802,7 +802,7 @@ impl Cheatcodes { } if let Some(result) = - self.strategy.inner.clone().zksync_try_create(self, ecx, &input, executor) + self.strategy.runner.clone().zksync_try_create(self, ecx, &input, executor) { return Some(result); } @@ -917,7 +917,7 @@ where { } } - self.strategy.inner.zksync_record_create_address(self.strategy.context.as_mut(), &outcome); + self.strategy.runner.zksync_record_create_address(self.strategy.context.as_mut(), &outcome); outcome } @@ -960,7 +960,7 @@ where { let prev = account.info.nonce; let nonce = prev.saturating_sub(1); account.info.nonce = nonce; - self.strategy.inner.zksync_sync_nonce( + self.strategy.runner.zksync_sync_nonce( self.strategy.context.as_mut(), sender, nonce, @@ -999,7 +999,7 @@ where { return None; } - self.strategy.inner.zksync_set_deployer_call_input(self.strategy.context.as_mut(), call); + self.strategy.runner.zksync_set_deployer_call_input(self.strategy.context.as_mut(), call); // Handle expected calls @@ -1133,7 +1133,7 @@ where { }) } - self.strategy.inner.record_broadcastable_call_transactions( + self.strategy.runner.record_broadcastable_call_transactions( self.strategy.context.as_mut(), self.config.clone(), call, @@ -1214,7 +1214,8 @@ where { }]); } - if let Some(result) = self.strategy.inner.clone().zksync_try_call(self, ecx, call, executor) + if let Some(result) = + self.strategy.runner.clone().zksync_try_call(self, ecx, call, executor) { return Some(result); } @@ -1276,7 +1277,7 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes { self.gas_metering.paused_frames.push(interpreter.gas); } - self.strategy.inner.post_initialize_interp( + self.strategy.runner.post_initialize_interp( self.strategy.context.as_mut(), interpreter, ecx, @@ -1325,7 +1326,7 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes { #[inline] fn step_end(&mut self, interpreter: &mut Interpreter, ecx: Ecx) { - if self.strategy.inner.pre_step_end(self.strategy.context.as_mut(), interpreter, ecx) { + if self.strategy.runner.pre_step_end(self.strategy.context.as_mut(), interpreter, ecx) { return; } diff --git a/crates/cheatcodes/src/strategy.rs b/crates/cheatcodes/src/strategy.rs index dc1928b40..42d1f6fbf 100644 --- a/crates/cheatcodes/src/strategy.rs +++ b/crates/cheatcodes/src/strategy.rs @@ -18,15 +18,19 @@ use crate::{ CheatsConfig, CheatsCtxt, Result, }; +/// Represents the context for [CheatcodeInspectorStrategy]. pub trait CheatcodeInspectorStrategyContext: Debug + Send + Sync + Any { + /// Clone the strategy context. fn new_cloned(&self) -> Box; - fn as_any_mut(&mut self) -> &mut dyn Any; + /// Alias as immutable reference of [Any]. fn as_any_ref(&self) -> &dyn Any; + /// Alias as mutable reference of [Any]. + fn as_any_mut(&mut self) -> &mut dyn Any; } impl CheatcodeInspectorStrategyContext for () { fn new_cloned(&self) -> Box { - Box::new(*self) + Box::new(()) } fn as_any_mut(&mut self) -> &mut dyn Any { @@ -38,26 +42,36 @@ impl CheatcodeInspectorStrategyContext for () { } } +/// Represents the strategy. #[derive(Debug)] -pub struct Strategy { - pub inner: Box, +pub struct CheatcodeInspectorStrategy { + /// Strategy runner. + pub runner: Box, + /// Strategy context. pub context: Box, } -pub fn new_evm_strategy() -> Strategy { - Strategy { inner: Box::new(EvmCheatcodeInspectorStrategy::default()), context: Box::new(()) } +impl CheatcodeInspectorStrategy { + pub fn new_evm() -> Self { + Self { + runner: Box::new(EvmCheatcodeInspectorStrategyRunner::default()), + context: Box::new(()), + } + } } -impl Clone for Strategy { +impl Clone for CheatcodeInspectorStrategy { fn clone(&self) -> Self { - Self { inner: self.inner.new_cloned(), context: self.context.new_cloned() } + Self { runner: self.runner.new_cloned(), context: self.context.new_cloned() } } } -pub trait CheatcodeInspectorStrategy: Debug + Send + Sync + CheatcodeInspectorStrategyExt { +pub trait CheatcodeInspectorStrategyRunner: + Debug + Send + Sync + CheatcodeInspectorStrategyExt +{ fn name(&self) -> &'static str; - fn new_cloned(&self) -> Box; + fn new_cloned(&self) -> Box; /// Get nonce. fn get_nonce(&self, ccx: &mut CheatsCtxt, address: Address) -> Result { @@ -328,14 +342,14 @@ pub trait CheatcodeInspectorStrategyExt { } #[derive(Debug, Default, Clone)] -pub struct EvmCheatcodeInspectorStrategy {} +pub struct EvmCheatcodeInspectorStrategyRunner {} -impl CheatcodeInspectorStrategy for EvmCheatcodeInspectorStrategy { +impl CheatcodeInspectorStrategyRunner for EvmCheatcodeInspectorStrategyRunner { fn name(&self) -> &'static str { "evm" } - fn new_cloned(&self) -> Box { + fn new_cloned(&self) -> Box { Box::new(self.clone()) } @@ -411,13 +425,13 @@ impl CheatcodeInspectorStrategy for EvmCheatcodeInspectorStrategy { } } -impl CheatcodeInspectorStrategyExt for EvmCheatcodeInspectorStrategy {} +impl CheatcodeInspectorStrategyExt for EvmCheatcodeInspectorStrategyRunner {} -impl Clone for Box { +impl Clone for Box { fn clone(&self) -> Self { self.new_cloned() } } -struct _ObjectSafe0(dyn CheatcodeInspectorStrategy); +struct _ObjectSafe0(dyn CheatcodeInspectorStrategyRunner); struct _ObjectSafe1(dyn CheatcodeInspectorStrategyExt); diff --git a/crates/cheatcodes/src/test.rs b/crates/cheatcodes/src/test.rs index 9c2e811b3..ae59439fc 100644 --- a/crates/cheatcodes/src/test.rs +++ b/crates/cheatcodes/src/test.rs @@ -15,7 +15,7 @@ impl Cheatcode for zkVmCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { enable } = *self; - ccx.state.strategy.inner.zksync_cheatcode_select_zk_vm( + ccx.state.strategy.runner.zksync_cheatcode_select_zk_vm( ccx.state.strategy.context.as_mut(), ccx.ecx, enable, @@ -27,14 +27,14 @@ impl Cheatcode for zkVmCall { impl Cheatcode for zkVmSkipCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { - ccx.state.strategy.inner.zksync_cheatcode_skip_zkvm(ccx.state.strategy.context.as_mut()) + ccx.state.strategy.runner.zksync_cheatcode_skip_zkvm(ccx.state.strategy.context.as_mut()) } } impl Cheatcode for zkUsePaymasterCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { paymaster_address, paymaster_input } = self; - ccx.state.strategy.inner.zksync_cheatcode_set_paymaster( + ccx.state.strategy.runner.zksync_cheatcode_set_paymaster( ccx.state.strategy.context.as_mut(), *paymaster_address, paymaster_input, @@ -47,7 +47,7 @@ impl Cheatcode for zkUseFactoryDepCall { let Self { name } = self; ccx.state .strategy - .inner + .runner .zksync_cheatcode_use_factory_deps(ccx.state.strategy.context.as_mut(), name.clone()) } } @@ -63,7 +63,7 @@ impl Cheatcode for zkRegisterContractCall { zkDeployedBytecode, } = self; - ccx.state.strategy.inner.zksync_cheatcode_register_contract( + ccx.state.strategy.runner.zksync_cheatcode_register_contract( ccx.state.strategy.context.as_mut(), name.clone(), zkBytecodeHash.0.into(), diff --git a/crates/chisel/src/executor.rs b/crates/chisel/src/executor.rs index 285d55496..4ae9a0e18 100644 --- a/crates/chisel/src/executor.rs +++ b/crates/chisel/src/executor.rs @@ -323,7 +323,7 @@ impl SessionSource { Some(backend) => backend, None => { let fork = self.config.evm_opts.get_fork(&self.config.foundry_config, env.clone()); - let backend = Backend::spawn(fork, strategy.inner.new_backend_strategy()); + let backend = Backend::spawn(fork, strategy.runner.new_backend_strategy()); self.config.backend = Some(backend.clone()); backend } @@ -339,7 +339,7 @@ impl SessionSource { None, None, Some(self.solc.version.clone()), - strategy.inner.new_cheatcode_inspector_strategy(strategy.context.as_mut()), + strategy.runner.new_cheatcode_inspector_strategy(strategy.context.as_mut()), ) .into(), ) diff --git a/crates/cli/src/utils/mod.rs b/crates/cli/src/utils/mod.rs index e06ff934a..c7025e5d9 100644 --- a/crates/cli/src/utils/mod.rs +++ b/crates/cli/src/utils/mod.rs @@ -9,8 +9,8 @@ use foundry_common::{ shell, }; use foundry_config::{Chain, Config}; -use foundry_evm::executors::strategy::{new_evm_strategy, Strategy}; -use foundry_strategy_zksync::new_zkysnc_strategy; +use foundry_evm::executors::strategy::ExecutorStrategy; +use foundry_strategy_zksync::ZksyncExecutorStrategyBuilder; use serde::de::DeserializeOwned; use std::{ ffi::OsStr, @@ -93,13 +93,13 @@ pub fn get_provider(config: &Config) -> Result { get_provider_builder(config)?.build() } -pub fn get_executor_strategy(config: &Config) -> Strategy { +pub fn get_executor_strategy(config: &Config) -> ExecutorStrategy { if config.zksync.should_compile() { info!("using zksync strategy"); - new_zkysnc_strategy() + ExecutorStrategy::new_zksync() } else { info!("using evm strategy"); - new_evm_strategy() + ExecutorStrategy::new_evm() } } diff --git a/crates/evm/core/src/backend/cow.rs b/crates/evm/core/src/backend/cow.rs index f53b6ce15..1a38ec6e8 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::{strategy::Strategy, BackendError, ForkInfo}; +use super::{strategy::BackendStrategy, BackendError, ForkInfo}; use crate::{ backend::{ diagnostic::RevertDiagnostic, Backend, DatabaseExt, LocalForkId, RevertStateSnapshotAction, @@ -92,7 +92,7 @@ impl DatabaseExt for CowBackend<'_> { self.backend.to_mut().get_fork_info(id) } - fn get_strategy(&mut self) -> &mut Strategy { + fn get_strategy(&mut self) -> &mut BackendStrategy { &mut self.backend.to_mut().strategy } diff --git a/crates/evm/core/src/backend/mod.rs b/crates/evm/core/src/backend/mod.rs index 4078e7736..b712f7c70 100644 --- a/crates/evm/core/src/backend/mod.rs +++ b/crates/evm/core/src/backend/mod.rs @@ -30,7 +30,7 @@ use std::{ collections::{BTreeMap, HashSet}, time::Instant, }; -use strategy::{BackendStrategyForkInfo, Strategy}; +use strategy::{BackendStrategy, BackendStrategyForkInfo}; mod diagnostic; pub use diagnostic::RevertDiagnostic; @@ -101,7 +101,7 @@ pub trait DatabaseExt: Database + DatabaseCommit { /// and the the fork environment. fn get_fork_info(&mut self, id: LocalForkId) -> eyre::Result; - fn get_strategy(&mut self) -> &mut Strategy; + fn get_strategy(&mut self) -> &mut BackendStrategy; /// Reverts the snapshot if it exists /// @@ -459,7 +459,7 @@ struct _ObjectSafe(dyn DatabaseExt); #[must_use] pub struct Backend { /// The behavior strategy. - pub strategy: Strategy, + pub strategy: BackendStrategy, /// The access point for managing forks forks: MultiFork, @@ -511,7 +511,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: Strategy) -> Self { + pub fn spawn(fork: Option, strategy: BackendStrategy) -> Self { Self::new(MultiFork::spawn(), fork, strategy) } @@ -521,7 +521,7 @@ impl Backend { /// database. /// /// Prefer using [`spawn`](Self::spawn) instead. - pub fn new(forks: MultiFork, fork: Option, strategy: Strategy) -> Self { + pub fn new(forks: MultiFork, fork: Option, strategy: BackendStrategy) -> Self { trace!(target: "backend", forking_mode=?fork.is_some(), "creating executor backend"); // Note: this will take of registering the `fork` let inner = BackendInner { @@ -560,7 +560,7 @@ impl Backend { /// Creates a new instance of `Backend` with fork added to the fork database and sets the fork /// as active pub(crate) fn new_with_fork( - strategy: Strategy, + strategy: BackendStrategy, id: &ForkId, fork: Fork, journaled_state: JournaledState, @@ -933,7 +933,7 @@ impl DatabaseExt for Backend { Ok(ForkInfo { fork_type, fork_env }) } - fn get_strategy(&mut self) -> &mut Strategy { + fn get_strategy(&mut self) -> &mut BackendStrategy { &mut self.strategy } @@ -1166,7 +1166,7 @@ impl DatabaseExt for Backend { let active_fork = self.active_fork_ids.map(|(_, idx)| self.inner.get_fork(idx)); // let active_fork = self.active_fork().cloned(); - self.strategy.inner.update_fork_db( + self.strategy.runner.update_fork_db( self.strategy.context.as_mut(), BackendStrategyForkInfo { active_fork, @@ -1226,7 +1226,7 @@ impl DatabaseExt for Backend { active.journaled_state.depth = journaled_state.depth; for addr in persistent_addrs { - self.strategy.inner.merge_journaled_state_data( + self.strategy.runner.merge_journaled_state_data( self.strategy.context.as_mut(), addr, journaled_state, @@ -1244,7 +1244,7 @@ impl DatabaseExt for Backend { for (addr, acc) in journaled_state.state.iter() { if acc.is_created() { if acc.is_touched() { - self.strategy.inner.merge_journaled_state_data( + self.strategy.runner.merge_journaled_state_data( self.strategy.context.as_mut(), *addr, journaled_state, @@ -1806,7 +1806,7 @@ impl BackendInner { pub fn roll_fork( &mut self, - strategy: &mut Strategy, + strategy: &mut BackendStrategy, id: LocalForkId, new_fork_id: ForkId, backend: SharedBackend, @@ -1818,7 +1818,7 @@ impl BackendInner { // we initialize a _new_ `ForkDB` but keep the state of persistent accounts let mut new_db = ForkDB::new(backend); for addr in self.persistent_accounts.iter().copied() { - strategy.inner.merge_db_account_data( + strategy.runner.merge_db_account_data( strategy.context.as_mut(), addr, &active.db, @@ -1936,7 +1936,7 @@ fn update_env_block(env: &mut Env, block: &AnyRpcBlock) { /// state, with an inspector. #[allow(clippy::too_many_arguments)] fn commit_transaction( - strategy: &mut Strategy, + strategy: &mut BackendStrategy, tx: &Transaction, mut env: EnvWithHandlerCfg, journaled_state: &mut JournaledState, @@ -2011,7 +2011,7 @@ fn apply_state_changeset( #[allow(clippy::needless_return)] mod tests { use crate::{ - backend::{strategy::new_evm_strategy, Backend}, + backend::{strategy::BackendStrategy, Backend}, fork::CreateFork, opts::EvmOpts, }; @@ -2045,7 +2045,7 @@ mod tests { evm_opts, }; - let backend = Backend::spawn(Some(fork), new_evm_strategy()); + let backend = Backend::spawn(Some(fork), BackendStrategy::new_evm()); // some rng contract from etherscan let address: Address = "63091244180ae240c87d1f528f5f269134cb07b3".parse().unwrap(); diff --git a/crates/evm/core/src/backend/strategy.rs b/crates/evm/core/src/backend/strategy.rs index b54eab781..1cdf78f33 100644 --- a/crates/evm/core/src/backend/strategy.rs +++ b/crates/evm/core/src/backend/strategy.rs @@ -11,9 +11,13 @@ pub struct BackendStrategyForkInfo<'a> { pub target_type: ForkType, } +/// Context for [BackendStrategyRunner]. pub trait BackendStrategyContext: Debug + Send + Sync + Any { + /// Clone the strategy context. fn new_cloned(&self) -> Box; + /// Alias as immutable reference of [Any]. fn as_any_ref(&self) -> &dyn Any; + /// Alias as mutable reference of [Any]. fn as_any_mut(&mut self) -> &mut dyn Any; } @@ -31,26 +35,32 @@ impl BackendStrategyContext for () { } } +/// Strategy for [super::Backend]. #[derive(Debug)] -pub struct Strategy { - pub inner: Box, +pub struct BackendStrategy { + /// Strategy runner. + pub runner: Box, + /// Strategy context. pub context: Box, } -pub fn new_evm_strategy() -> Strategy { - Strategy { inner: Box::new(EvmBackendStrategy), context: Box::new(()) } +impl BackendStrategy { + /// Create a new instance of [BackendStrategy] + pub fn new_evm() -> Self { + Self { runner: Box::new(EvmBackendStrategyRunner), context: Box::new(()) } + } } -impl Clone for Strategy { +impl Clone for BackendStrategy { fn clone(&self) -> Self { - Self { inner: self.inner.new_cloned(), context: self.context.new_cloned() } + Self { runner: self.runner.new_cloned(), context: self.context.new_cloned() } } } -pub trait BackendStrategy: Debug + Send + Sync + BackendStrategyExt { +pub trait BackendStrategyRunner: Debug + Send + Sync + BackendStrategyRunnerExt { fn name(&self) -> &'static str; - fn new_cloned(&self) -> Box; + fn new_cloned(&self) -> Box; /// When creating or switching forks, we update the AccountInfo of the contract fn update_fork_db( @@ -81,7 +91,7 @@ pub trait BackendStrategy: Debug + Send + Sync + BackendStrategyExt { ); } -pub trait BackendStrategyExt { +pub trait BackendStrategyRunnerExt { /// Saves the storage keys for immutable variables per address. /// /// These are required during fork to help merge the persisted addresses, as they are stored @@ -97,17 +107,17 @@ pub trait BackendStrategyExt { } } -struct _ObjectSafe(dyn BackendStrategy); +struct _ObjectSafe(dyn BackendStrategyRunner); #[derive(Debug, Default, Clone, Serialize, Deserialize)] -pub struct EvmBackendStrategy; +pub struct EvmBackendStrategyRunner; -impl BackendStrategy for EvmBackendStrategy { +impl BackendStrategyRunner for EvmBackendStrategyRunner { fn name(&self) -> &'static str { "evm" } - fn new_cloned(&self) -> Box { + fn new_cloned(&self) -> Box { Box::new(self.clone()) } @@ -154,9 +164,9 @@ impl BackendStrategy for EvmBackendStrategy { } } -impl BackendStrategyExt for EvmBackendStrategy {} +impl BackendStrategyRunnerExt for EvmBackendStrategyRunner {} -impl EvmBackendStrategy { +impl EvmBackendStrategyRunner { /// Merges the state of all `accounts` from the currently active db into the given `fork` pub(crate) fn update_fork_db_contracts( &self, @@ -258,7 +268,7 @@ impl EvmBackendMergeStrategy { } } -impl Clone for Box { +impl Clone for Box { fn clone(&self) -> Self { self.new_cloned() } diff --git a/crates/evm/evm/src/executors/builder.rs b/crates/evm/evm/src/executors/builder.rs index 3dc7009df..20ac319bb 100644 --- a/crates/evm/evm/src/executors/builder.rs +++ b/crates/evm/evm/src/executors/builder.rs @@ -2,7 +2,7 @@ use crate::{executors::Executor, inspectors::InspectorStackBuilder}; use foundry_evm_core::backend::Backend; use revm::primitives::{Env, EnvWithHandlerCfg, SpecId}; -use super::strategy::Strategy; +use super::strategy::ExecutorStrategy; /// The builder that allows to configure an evm [`Executor`] which a stack of optional /// [`revm::Inspector`]s, such as [`Cheatcodes`]. @@ -76,7 +76,7 @@ impl ExecutorBuilder { /// Builds the executor as configured. #[inline] - pub fn build(self, env: Env, db: Backend, strategy: Strategy) -> Executor { + pub fn build(self, env: Env, db: Backend, strategy: ExecutorStrategy) -> Executor { let Self { mut stack, gas_limit, spec_id, legacy_assertions } = self; if stack.block.is_none() { stack.block = Some(env.block.clone()); diff --git a/crates/evm/evm/src/executors/mod.rs b/crates/evm/evm/src/executors/mod.rs index 352148fc2..061f2ac11 100644 --- a/crates/evm/evm/src/executors/mod.rs +++ b/crates/evm/evm/src/executors/mod.rs @@ -37,7 +37,7 @@ use revm::{ }, }; use std::borrow::Cow; -use strategy::Strategy; +use strategy::ExecutorStrategy; mod builder; pub use builder::ExecutorBuilder; @@ -93,7 +93,7 @@ pub struct Executor { /// Whether `failed()` should be called on the test contract to determine if the test failed. legacy_assertions: bool, - strategy: Strategy, + strategy: ExecutorStrategy, } impl Clone for Executor { @@ -124,7 +124,7 @@ impl Executor { inspector: InspectorStack, gas_limit: u64, legacy_assertions: bool, - strategy: Strategy, + strategy: ExecutorStrategy, ) -> Self { // Need to create a non-empty contract on the cheatcodes address so `extcodesize` checks // do not fail. @@ -217,7 +217,7 @@ impl Executor { /// Set the balance of an account. pub fn set_balance(&mut self, address: Address, amount: U256) -> BackendResult<()> { trace!(?address, ?amount, "setting account balance"); - self.strategy.inner.clone().set_balance(self, address, amount) + self.strategy.runner.clone().set_balance(self, address, amount) } /// Gets the balance of an account @@ -227,7 +227,7 @@ impl Executor { /// Set the nonce of an account. pub fn set_nonce(&mut self, address: Address, nonce: u64) -> BackendResult<()> { - self.strategy.inner.clone().set_nonce(self, address, nonce) + self.strategy.runner.clone().set_nonce(self, address, nonce) } /// Returns the nonce of an account. @@ -260,7 +260,7 @@ impl Executor { #[inline] pub fn set_transaction_other_fields(&mut self, other_fields: OtherFields) { - self.strategy.inner.set_inspect_context(self.strategy.context.as_mut(), other_fields); + self.strategy.runner.set_inspect_context(self.strategy.context.as_mut(), other_fields); } /// Deploys a contract and commits the new state to the underlying database. @@ -443,7 +443,7 @@ impl Executor { backend.is_initialized = false; backend.spec_id = env.spec_id(); - let result = self.strategy.inner.call_inspect( + let result = self.strategy.runner.call_inspect( self.strategy.context.as_ref(), &mut backend, &mut env, @@ -465,7 +465,7 @@ impl Executor { let backend = &mut self.backend; backend.initialize(&env); - let result_and_state = self.strategy.inner.transact_inspect( + let result_and_state = self.strategy.runner.transact_inspect( self.strategy.context.as_mut(), backend, &mut env, diff --git a/crates/evm/evm/src/executors/strategy.rs b/crates/evm/evm/src/executors/strategy.rs index 2d5fcfd84..2a88deb3b 100644 --- a/crates/evm/evm/src/executors/strategy.rs +++ b/crates/evm/evm/src/executors/strategy.rs @@ -3,9 +3,11 @@ use std::{any::Any, fmt::Debug}; use alloy_primitives::{Address, U256}; use alloy_serde::OtherFields; use eyre::{Context, Result}; -use foundry_cheatcodes::strategy::EvmCheatcodeInspectorStrategy; +use foundry_cheatcodes::strategy::{ + CheatcodeInspectorStrategy, EvmCheatcodeInspectorStrategyRunner, +}; use foundry_evm_core::{ - backend::{BackendResult, DatabaseExt}, + backend::{strategy::BackendStrategy, BackendResult, DatabaseExt}, InspectorExt, }; use foundry_zksync_compiler::DualCompiledContracts; @@ -17,8 +19,11 @@ use revm::{ use super::Executor; pub trait ExecutorStrategyContext: Debug + Send + Sync + Any { + /// Clone the strategy context. fn new_cloned(&self) -> Box; + /// Alias as immutable reference of [Any]. fn as_any_ref(&self) -> &dyn Any; + /// Alias as mutable reference of [Any]. fn as_any_mut(&mut self) -> &mut dyn Any; } @@ -37,25 +42,29 @@ impl ExecutorStrategyContext for () { } #[derive(Debug)] -pub struct Strategy { - pub inner: Box, +pub struct ExecutorStrategy { + /// Strategy runner. + pub runner: Box, + /// Strategy context. pub context: Box, } -pub fn new_evm_strategy() -> Strategy { - Strategy { inner: Box::new(EvmExecutorStrategy::default()), context: Box::new(()) } +impl ExecutorStrategy { + pub fn new_evm() -> Self { + Self { runner: Box::new(EvmExecutorStrategyRunner::default()), context: Box::new(()) } + } } -impl Clone for Strategy { +impl Clone for ExecutorStrategy { fn clone(&self) -> Self { - Self { inner: self.inner.new_cloned(), context: self.context.new_cloned() } + Self { runner: self.runner.new_cloned(), context: self.context.new_cloned() } } } -pub trait ExecutorStrategy: Debug + Send + Sync + ExecutorStrategyExt { +pub trait ExecutorStrategyRunner: Debug + Send + Sync + ExecutorStrategyExt { fn name(&self) -> &'static str; - fn new_cloned(&self) -> Box; + fn new_cloned(&self) -> Box; fn set_balance( &self, @@ -86,16 +95,18 @@ pub trait ExecutorStrategy: Debug + Send + Sync + ExecutorStrategyExt { inspector: &mut dyn InspectorExt, ) -> eyre::Result; - fn new_backend_strategy(&self) -> foundry_evm_core::backend::strategy::Strategy; + fn new_backend_strategy(&self) -> BackendStrategy; fn new_cheatcode_inspector_strategy( &self, ctx: &dyn ExecutorStrategyContext, - ) -> foundry_cheatcodes::strategy::Strategy; + ) -> foundry_cheatcodes::strategy::CheatcodeInspectorStrategy; // TODO perhaps need to create fresh strategies as well } +/// Extended trait for ZKsync. pub trait ExecutorStrategyExt { + /// Set [DualCompiledContracts] on the context. fn zksync_set_dual_compiled_contracts( &self, _ctx: &mut dyn ExecutorStrategyContext, @@ -103,6 +114,7 @@ pub trait ExecutorStrategyExt { ) { } + /// Set the fork environment on the context. fn zksync_set_fork_env( &self, _ctx: &mut dyn ExecutorStrategyContext, @@ -113,15 +125,16 @@ pub trait ExecutorStrategyExt { } } +/// Implements [ExecutorStrategyRunner] for EVM. #[derive(Debug, Default, Clone)] -pub struct EvmExecutorStrategy {} +pub struct EvmExecutorStrategyRunner {} -impl ExecutorStrategy for EvmExecutorStrategy { +impl ExecutorStrategyRunner for EvmExecutorStrategyRunner { fn name(&self) -> &'static str { "evm" } - fn new_cloned(&self) -> Box { + fn new_cloned(&self) -> Box { Box::new(self.clone()) } @@ -201,24 +214,24 @@ impl ExecutorStrategy for EvmExecutorStrategy { Ok(()) } - fn new_backend_strategy(&self) -> foundry_evm_core::backend::strategy::Strategy { - foundry_evm_core::backend::strategy::new_evm_strategy() + fn new_backend_strategy(&self) -> BackendStrategy { + BackendStrategy::new_evm() } fn new_cheatcode_inspector_strategy( &self, _ctx: &dyn ExecutorStrategyContext, - ) -> foundry_cheatcodes::strategy::Strategy { - foundry_cheatcodes::strategy::Strategy { - inner: Box::new(EvmCheatcodeInspectorStrategy::default()), + ) -> CheatcodeInspectorStrategy { + CheatcodeInspectorStrategy { + runner: Box::new(EvmCheatcodeInspectorStrategyRunner::default()), context: Box::new(()), } } } -impl ExecutorStrategyExt for EvmExecutorStrategy {} +impl ExecutorStrategyExt for EvmExecutorStrategyRunner {} -impl Clone for Box { +impl Clone for Box { fn clone(&self) -> Self { self.new_cloned() } diff --git a/crates/evm/evm/src/executors/trace.rs b/crates/evm/evm/src/executors/trace.rs index be0be748e..1b523897d 100644 --- a/crates/evm/evm/src/executors/trace.rs +++ b/crates/evm/evm/src/executors/trace.rs @@ -6,7 +6,7 @@ use foundry_evm_traces::{InternalTraceMode, TraceMode}; use revm::primitives::{Env, SpecId}; use std::ops::{Deref, DerefMut}; -use super::strategy::Strategy; +use super::strategy::ExecutorStrategy; /// A default executor with tracing enabled pub struct TracingExecutor { @@ -21,9 +21,9 @@ impl TracingExecutor { debug: bool, decode_internal: bool, alphanet: bool, - strategy: Strategy, + strategy: ExecutorStrategy, ) -> Self { - let db = Backend::spawn(fork, strategy.inner.new_backend_strategy()); + let db = Backend::spawn(fork, strategy.runner.new_backend_strategy()); let trace_mode = TraceMode::Call.with_debug(debug).with_decode_internal(if decode_internal { InternalTraceMode::Full diff --git a/crates/forge/bin/cmd/test/mod.rs b/crates/forge/bin/cmd/test/mod.rs index 96335bdc7..425c82598 100644 --- a/crates/forge/bin/cmd/test/mod.rs +++ b/crates/forge/bin/cmd/test/mod.rs @@ -368,7 +368,7 @@ impl TestArgs { // Prepare the test builder. let config = Arc::new(config); - strategy.inner.zksync_set_dual_compiled_contracts( + strategy.runner.zksync_set_dual_compiled_contracts( strategy.context.as_mut(), dual_compiled_contracts.unwrap_or_default(), ); diff --git a/crates/forge/src/multi_runner.rs b/crates/forge/src/multi_runner.rs index c716318da..0c9c7731b 100644 --- a/crates/forge/src/multi_runner.rs +++ b/crates/forge/src/multi_runner.rs @@ -18,7 +18,7 @@ use foundry_config::Config; use foundry_evm::{ backend::Backend, decode::RevertDecoder, - executors::{strategy::Strategy, ExecutorBuilder}, + executors::{strategy::ExecutorStrategy, ExecutorBuilder}, fork::CreateFork, inspectors::CheatsConfig, opts::EvmOpts, @@ -85,7 +85,7 @@ pub struct MultiContractRunner { /// Library addresses used to link contracts. pub libraries: Libraries, /// Execution strategy. - pub strategy: Strategy, + pub strategy: ExecutorStrategy, } impl MultiContractRunner { @@ -178,7 +178,7 @@ impl MultiContractRunner { trace!("running all tests"); // The DB backend that serves all the data. - let db = Backend::spawn(self.fork.take(), self.strategy.inner.new_backend_strategy()); + let db = Backend::spawn(self.fork.take(), self.strategy.runner.new_backend_strategy()); let find_timer = Instant::now(); let contracts = self.matching_contracts(filter).collect::>(); @@ -250,7 +250,7 @@ impl MultiContractRunner { Some(self.known_contracts.clone()), Some(artifact_id.name.clone()), Some(artifact_id.version.clone()), - self.strategy.inner.new_cheatcode_inspector_strategy(self.strategy.context.as_ref()), + self.strategy.runner.new_cheatcode_inspector_strategy(self.strategy.context.as_ref()), ); let trace_mode = TraceMode::default() @@ -407,7 +407,7 @@ impl MultiContractRunnerBuilder { zk_output: Option, env: revm::primitives::Env, evm_opts: EvmOpts, - strategy: Strategy, + strategy: ExecutorStrategy, ) -> Result { let mut known_contracts = ContractsByArtifact::default(); let output = output.with_stripped_file_prefixes(root); diff --git a/crates/forge/src/runner.rs b/crates/forge/src/runner.rs index 5f29c02d4..745e0c2a2 100644 --- a/crates/forge/src/runner.rs +++ b/crates/forge/src/runner.rs @@ -151,7 +151,7 @@ impl ContractRunner<'_> { // nonce. if let Some(cheatcodes) = &mut self.executor.inspector.cheatcodes { debug!("test contract deployed"); - cheatcodes.strategy.inner.base_contract_deployed(cheatcodes.strategy.context.as_mut()); + cheatcodes.strategy.runner.base_contract_deployed(cheatcodes.strategy.context.as_mut()); } // Optionally call the `setUp` function diff --git a/crates/forge/tests/it/test_helpers.rs b/crates/forge/tests/it/test_helpers.rs index beb7e4f88..3446b8ec1 100644 --- a/crates/forge/tests/it/test_helpers.rs +++ b/crates/forge/tests/it/test_helpers.rs @@ -361,7 +361,7 @@ impl ForgeTestData { let mut strategy = utils::get_executor_strategy(&zk_config); strategy - .inner + .runner .zksync_set_dual_compiled_contracts(strategy.context.as_mut(), dual_compiled_contracts); let mut builder = self.base_runner(); builder.config = Arc::new(zk_config); diff --git a/crates/script/src/lib.rs b/crates/script/src/lib.rs index 80d437744..169aa7172 100644 --- a/crates/script/src/lib.rs +++ b/crates/script/src/lib.rs @@ -598,7 +598,7 @@ impl ScriptConfig { Some(db) => db.clone(), None => { let fork = self.evm_opts.get_fork(&self.config, env.clone()); - let backend = Backend::spawn(fork, strategy.inner.new_backend_strategy()); + let backend = Backend::spawn(fork, strategy.runner.new_backend_strategy()); self.backends.insert(fork_url.clone(), backend.clone()); backend } @@ -607,7 +607,7 @@ impl ScriptConfig { // It's only really `None`, when we don't pass any `--fork-url`. And if so, there is // no need to cache it, since there won't be any onchain simulation that we'd need // to cache the backend for. - Backend::spawn(None, strategy.inner.new_backend_strategy()) + Backend::spawn(None, strategy.runner.new_backend_strategy()) }; // We need to enable tracing to decode contract names: local or external. @@ -624,13 +624,13 @@ impl ScriptConfig { if let Some((known_contracts, script_wallets, target, dual_compiled_contracts)) = cheats_data { - strategy.inner.zksync_set_dual_compiled_contracts( + strategy.runner.zksync_set_dual_compiled_contracts( strategy.context.as_mut(), dual_compiled_contracts, ); if let Some(fork_url) = &self.evm_opts.fork_url { - strategy.inner.zksync_set_fork_env(strategy.context.as_mut(), fork_url, &env)?; + strategy.runner.zksync_set_fork_env(strategy.context.as_mut(), fork_url, &env)?; } builder = builder.inspectors(|stack| { @@ -643,7 +643,7 @@ impl ScriptConfig { Some(target.name), Some(target.version), strategy - .inner + .runner .new_cheatcode_inspector_strategy(strategy.context.as_ref()), ) .into(), diff --git a/crates/script/src/runner.rs b/crates/script/src/runner.rs index cb87bffc8..22b26446d 100644 --- a/crates/script/src/runner.rs +++ b/crates/script/src/runner.rs @@ -174,7 +174,7 @@ impl ScriptRunner { // nonce. if let Some(cheatcodes) = &mut self.executor.inspector.cheatcodes { debug!("script deployed"); - cheatcodes.strategy.inner.base_contract_deployed(cheatcodes.strategy.context.as_mut()); + cheatcodes.strategy.runner.base_contract_deployed(cheatcodes.strategy.context.as_mut()); } // Optionally call the `setUp` function diff --git a/crates/strategy/zksync/src/backend.rs b/crates/strategy/zksync/src/backend.rs index 9ab91be12..76a9286ea 100644 --- a/crates/strategy/zksync/src/backend.rs +++ b/crates/strategy/zksync/src/backend.rs @@ -1,10 +1,13 @@ use std::{any::Any, collections::hash_map::Entry}; use alloy_primitives::{map::HashMap, Address, U256}; -use foundry_evm::backend::strategy::{BackendStrategyContext, BackendStrategyExt}; +use foundry_evm::backend::strategy::{ + BackendStrategy, BackendStrategyContext, BackendStrategyRunnerExt, +}; use foundry_evm_core::backend::{ strategy::{ - BackendStrategy, BackendStrategyForkInfo, EvmBackendMergeStrategy, EvmBackendStrategy, + BackendStrategyForkInfo, BackendStrategyRunner, EvmBackendMergeStrategy, + EvmBackendStrategyRunner, }, BackendInner, Fork, ForkDB, FoundryEvmInMemoryDB, }; @@ -17,6 +20,7 @@ use revm::{db::CacheDB, primitives::HashSet, DatabaseRef, JournaledState}; use serde::{Deserialize, Serialize}; use tracing::trace; +/// Context for [ZksyncBackendStrategyRunner]. #[derive(Debug, Default, Clone, Serialize, Deserialize)] pub struct ZksyncBackendStrategyContext { /// Store storage keys per contract address for immutable variables. @@ -37,32 +41,18 @@ impl BackendStrategyContext for ZksyncBackendStrategyContext { } } +/// ZKsync implementation for [BackendStrategyRunner]. #[derive(Debug, Default, Clone, Serialize, Deserialize)] -pub struct ZksyncBackendStrategy { - evm: EvmBackendStrategy, -} - -#[derive(Debug, Default, Clone, Serialize, Deserialize)] -pub struct ZkBackendInspectData { - #[serde(skip_serializing_if = "Option::is_none")] - pub factory_deps: Option>>, - - #[serde(skip_serializing_if = "Option::is_none")] - pub paymaster_data: Option, - - pub use_evm: bool, -} - -fn get_context(ctx: &mut dyn BackendStrategyContext) -> &mut ZksyncBackendStrategyContext { - ctx.as_any_mut().downcast_mut().expect("expected ZksyncBackendStrategyContext") +pub struct ZksyncBackendStrategyRunner { + evm: EvmBackendStrategyRunner, } -impl BackendStrategy for ZksyncBackendStrategy { +impl BackendStrategyRunner for ZksyncBackendStrategyRunner { fn name(&self) -> &'static str { "zk" } - fn new_cloned(&self) -> Box { + fn new_cloned(&self) -> Box { Box::new(self.clone()) } @@ -126,7 +116,7 @@ impl BackendStrategy for ZksyncBackendStrategy { } } -impl ZksyncBackendStrategy { +impl ZksyncBackendStrategyRunner { /// Merges the state of all `accounts` from the currently active db into the given `fork` pub(crate) fn update_fork_db_contracts( &self, @@ -168,7 +158,7 @@ impl ZksyncBackendStrategy { } } -impl BackendStrategyExt for ZksyncBackendStrategy { +impl BackendStrategyRunnerExt for ZksyncBackendStrategyRunner { fn zksync_save_immutable_storage( &self, ctx: &mut dyn BackendStrategyContext, @@ -183,6 +173,21 @@ impl BackendStrategyExt for ZksyncBackendStrategy { } } +/// Create ZKsync strategy for [BackendStrategy]. +pub trait ZksyncBackendStrategyBuilder { + /// Create new zksync strategy. + fn new_zksync() -> Self; +} + +impl ZksyncBackendStrategyBuilder for BackendStrategy { + fn new_zksync() -> Self { + Self { + runner: Box::new(ZksyncBackendStrategyRunner::default()), + context: Box::new(ZksyncBackendStrategyContext::default()), + } + } +} + pub(crate) struct ZksyncBackendMerge; /// Defines the zksync specific state to help during merge. @@ -359,3 +364,7 @@ impl ZksyncBackendMerge { } } } + +fn get_context(ctx: &mut dyn BackendStrategyContext) -> &mut ZksyncBackendStrategyContext { + ctx.as_any_mut().downcast_mut().expect("expected ZksyncBackendStrategyContext") +} diff --git a/crates/strategy/zksync/src/cheatcode.rs b/crates/strategy/zksync/src/cheatcode.rs index f8efc80c5..08e6b886e 100644 --- a/crates/strategy/zksync/src/cheatcode.rs +++ b/crates/strategy/zksync/src/cheatcode.rs @@ -11,7 +11,8 @@ use foundry_cheatcodes::{ journaled_account, make_acc_non_empty, strategy::{ CheatcodeInspectorStrategy, CheatcodeInspectorStrategyContext, - CheatcodeInspectorStrategyExt, EvmCheatcodeInspectorStrategy, + CheatcodeInspectorStrategyExt, CheatcodeInspectorStrategyRunner, + EvmCheatcodeInspectorStrategyRunner, }, Broadcast, BroadcastableTransaction, BroadcastableTransactions, Cheatcodes, CheatcodesExecutor, CheatsConfig, CheatsCtxt, CommonCreateInput, DealRecord, Ecx, Error, InnerEcx, Result, Vm, @@ -81,11 +82,13 @@ macro_rules! bail { }; } +/// ZKsync implementation for [CheatcodeInspectorStrategyRunner]. #[derive(Debug, Default, Clone)] -pub struct ZksyncCheatcodeInspectorStrategy { - evm: EvmCheatcodeInspectorStrategy, +pub struct ZksyncCheatcodeInspectorStrategyRunner { + evm: EvmCheatcodeInspectorStrategyRunner, } +/// Context for [ZksyncCheatcodeInspectorStrategyRunner]. #[derive(Debug, Default, Clone)] pub struct ZksyncCheatcodeInspectorStrategyContext { pub using_zk_vm: bool, @@ -202,18 +205,6 @@ impl CheatcodeInspectorStrategyContext for ZksyncCheatcodeInspectorStrategyConte } } -fn get_context( - ctx: &mut dyn CheatcodeInspectorStrategyContext, -) -> &mut ZksyncCheatcodeInspectorStrategyContext { - ctx.as_any_mut().downcast_mut().expect("expected ZksyncCheatcodeInspectorStrategyContext") -} - -fn get_context_ref( - ctx: &dyn CheatcodeInspectorStrategyContext, -) -> &ZksyncCheatcodeInspectorStrategyContext { - ctx.as_any_ref().downcast_ref().expect("expected ZksyncCheatcodeInspectorStrategyContext") -} - /// Allows overriding nonce update behavior for the tx caller in the zkEVM. /// /// Since each CREATE or CALL is executed as a separate transaction within zkEVM, we currently skip @@ -247,12 +238,12 @@ impl ZkPersistNonceUpdate { } } -impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { +impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner { fn name(&self) -> &'static str { "zk" } - fn new_cloned(&self) -> Box { + fn new_cloned(&self) -> Box { Box::new(self.clone()) } @@ -753,7 +744,7 @@ impl CheatcodeInspectorStrategy for ZksyncCheatcodeInspectorStrategy { } } -impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { +impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategyRunner { fn zksync_cheatcode_skip_zkvm( &self, ctx: &mut dyn CheatcodeInspectorStrategyContext, @@ -1045,7 +1036,7 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { }) .collect::>(); let strategy = ecx.db.get_strategy(); - strategy.inner.zksync_save_immutable_storage( + strategy.runner.zksync_save_immutable_storage( strategy.context.as_mut(), addr, keys, @@ -1297,7 +1288,7 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategy { } } -impl ZksyncCheatcodeInspectorStrategy { +impl ZksyncCheatcodeInspectorStrategyRunner { /// Selects the appropriate VM for the fork. Options: EVM, ZK-VM. /// CALL and CREATE are handled by the selected VM. /// @@ -1506,6 +1497,58 @@ impl ZksyncCheatcodeInspectorStrategy { } } +/// Setting for migrating the database to zkEVM storage when starting in ZKsync mode. +/// The migration is performed on the DB via the inspector so must only be performed once. +#[derive(Debug, Default, Clone)] +pub enum ZkStartupMigration { + /// Defer database migration to a later execution point. + /// + /// This is required as we need to wait for some baseline deployments + /// to occur before the test/script execution is performed. + #[default] + Defer, + /// Allow database migration. + Allow, + /// Database migration has already been performed. + Done, +} + +impl ZkStartupMigration { + /// Check if startup migration is allowed. Migration is disallowed if it's to be deferred or has + /// already been performed. + pub fn is_allowed(&self) -> bool { + matches!(self, Self::Allow) + } + + /// Allow migrating the the DB to zkEVM storage. + pub fn allow(&mut self) { + *self = Self::Allow + } + + /// Mark the migration as completed. It must not be performed again. + pub fn done(&mut self) { + *self = Self::Done + } +} + +/// Create ZKsync strategy for [CheatcodeInspectorStrategy]. +pub trait ZksyncCheatcodeInspectorStrategyBuilder { + /// Create new ZKsync strategy. + fn new_zksync(dual_compiled_contracts: DualCompiledContracts, zk_env: ZkEnv) -> Self; +} + +impl ZksyncCheatcodeInspectorStrategyBuilder for CheatcodeInspectorStrategy { + fn new_zksync(dual_compiled_contracts: DualCompiledContracts, zk_env: ZkEnv) -> Self { + Self { + runner: Box::new(ZksyncCheatcodeInspectorStrategyRunner::default()), + context: Box::new(ZksyncCheatcodeInspectorStrategyContext::new( + dual_compiled_contracts, + zk_env, + )), + } + } +} + fn get_artifact_code( dual_compiled_contracts: &DualCompiledContracts, using_zk_vm: bool, @@ -1637,36 +1680,14 @@ fn get_artifact_code( maybe_bytecode.ok_or_else(|| fmt_err!("no bytecode for contract; is it abstract or unlinked?")) } -/// Setting for migrating the database to zkEVM storage when starting in ZKsync mode. -/// The migration is performed on the DB via the inspector so must only be performed once. -#[derive(Debug, Default, Clone)] -pub enum ZkStartupMigration { - /// Defer database migration to a later execution point. - /// - /// This is required as we need to wait for some baseline deployments - /// to occur before the test/script execution is performed. - #[default] - Defer, - /// Allow database migration. - Allow, - /// Database migration has already been performed. - Done, +fn get_context( + ctx: &mut dyn CheatcodeInspectorStrategyContext, +) -> &mut ZksyncCheatcodeInspectorStrategyContext { + ctx.as_any_mut().downcast_mut().expect("expected ZksyncCheatcodeInspectorStrategyContext") } -impl ZkStartupMigration { - /// Check if startup migration is allowed. Migration is disallowed if it's to be deferred or has - /// already been performed. - pub fn is_allowed(&self) -> bool { - matches!(self, Self::Allow) - } - - /// Allow migrating the the DB to zkEVM storage. - pub fn allow(&mut self) { - *self = Self::Allow - } - - /// Mark the migration as completed. It must not be performed again. - pub fn done(&mut self) { - *self = Self::Done - } +fn get_context_ref( + ctx: &dyn CheatcodeInspectorStrategyContext, +) -> &ZksyncCheatcodeInspectorStrategyContext { + ctx.as_any_ref().downcast_ref().expect("expected ZksyncCheatcodeInspectorStrategyContext") } diff --git a/crates/strategy/zksync/src/executor.rs b/crates/strategy/zksync/src/executor.rs index 824080792..191ec4145 100644 --- a/crates/strategy/zksync/src/executor.rs +++ b/crates/strategy/zksync/src/executor.rs @@ -7,7 +7,8 @@ use foundry_evm::{ backend::{BackendResult, DatabaseExt}, executors::{ strategy::{ - EvmExecutorStrategy, ExecutorStrategy, ExecutorStrategyContext, ExecutorStrategyExt, + EvmExecutorStrategyRunner, ExecutorStrategy, ExecutorStrategyContext, + ExecutorStrategyExt, ExecutorStrategyRunner, }, Executor, }, @@ -22,11 +23,11 @@ use revm::{ use zksync_types::H256; use crate::{ - backend::ZksyncBackendStrategyContext, - cheatcode::{ZksyncCheatcodeInspectorStrategyContext, ZKSYNC_TRANSACTION_OTHER_FIELDS_KEY}, - ZksyncBackendStrategy, ZksyncCheatcodeInspectorStrategy, + backend::ZksyncBackendStrategyBuilder, + cheatcode::{ZksyncCheatcodeInspectorStrategyBuilder, ZKSYNC_TRANSACTION_OTHER_FIELDS_KEY}, }; +/// Defines the context for [ZksyncExecutorStrategyRunner]. #[derive(Debug, Default, Clone)] pub struct ZksyncExecutorStrategyContext { inspect_context: Option, @@ -49,9 +50,10 @@ impl ExecutorStrategyContext for ZksyncExecutorStrategyContext { } } +/// Defines the [ExecutorStrategyRunner] strategy for ZKsync. #[derive(Debug, Default, Clone)] -pub struct ZksyncExecutorStrategy { - evm: EvmExecutorStrategy, +pub struct ZksyncExecutorStrategyRunner { + evm: EvmExecutorStrategyRunner, } fn get_context_ref(ctx: &dyn ExecutorStrategyContext) -> &ZksyncExecutorStrategyContext { @@ -62,12 +64,12 @@ fn get_context(ctx: &mut dyn ExecutorStrategyContext) -> &mut ZksyncExecutorStra ctx.as_any_mut().downcast_mut().expect("expected ZksyncExecutorStrategyContext") } -impl ExecutorStrategy for ZksyncExecutorStrategy { +impl ExecutorStrategyRunner for ZksyncExecutorStrategyRunner { fn name(&self) -> &'static str { "zk" } - fn new_cloned(&self) -> Box { + fn new_cloned(&self) -> Box { Box::new(self.clone()) } @@ -114,25 +116,19 @@ impl ExecutorStrategy for ZksyncExecutorStrategy { Ok(()) } - fn new_backend_strategy(&self) -> foundry_evm_core::backend::strategy::Strategy { - foundry_evm_core::backend::strategy::Strategy { - inner: Box::new(ZksyncBackendStrategy::default()), - context: Box::new(ZksyncBackendStrategyContext::default()), - } + fn new_backend_strategy(&self) -> foundry_evm_core::backend::strategy::BackendStrategy { + foundry_evm_core::backend::strategy::BackendStrategy::new_zksync() } fn new_cheatcode_inspector_strategy( &self, ctx: &dyn ExecutorStrategyContext, - ) -> foundry_cheatcodes::strategy::Strategy { + ) -> foundry_cheatcodes::strategy::CheatcodeInspectorStrategy { let ctx = get_context_ref(ctx); - foundry_cheatcodes::strategy::Strategy { - inner: Box::new(ZksyncCheatcodeInspectorStrategy::default()), - context: Box::new(ZksyncCheatcodeInspectorStrategyContext::new( - ctx.dual_compiled_contracts.clone(), - ctx.zk_env.clone(), - )), - } + foundry_cheatcodes::strategy::CheatcodeInspectorStrategy::new_zksync( + ctx.dual_compiled_contracts.clone(), + ctx.zk_env.clone(), + ) } fn call_inspect( @@ -187,7 +183,7 @@ impl ExecutorStrategy for ZksyncExecutorStrategy { } } -impl ExecutorStrategyExt for ZksyncExecutorStrategy { +impl ExecutorStrategyExt for ZksyncExecutorStrategyRunner { fn zksync_set_dual_compiled_contracts( &self, ctx: &mut dyn ExecutorStrategyContext, @@ -239,7 +235,7 @@ impl ExecutorStrategyExt for ZksyncExecutorStrategy { } } -/// Retrieve metadata for zksync tx +/// Retrieve metadata for ZKsync tx. pub fn get_zksync_transaction_metadata( other_fields: &OtherFields, ) -> Option { @@ -250,9 +246,17 @@ pub fn get_zksync_transaction_metadata( .flatten() } -pub fn new_zkysnc_strategy() -> foundry_evm::executors::strategy::Strategy { - foundry_evm::executors::strategy::Strategy { - inner: Box::new(ZksyncExecutorStrategy::default()), - context: Box::new(ZksyncExecutorStrategyContext::default()), +/// Create ZKsync strategy for [ExecutorStrategy]. +pub trait ZksyncExecutorStrategyBuilder { + /// Create new zksync strategy. + fn new_zksync() -> Self; +} + +impl ZksyncExecutorStrategyBuilder for ExecutorStrategy { + fn new_zksync() -> Self { + Self { + runner: Box::new(ZksyncExecutorStrategyRunner::default()), + context: Box::new(ZksyncExecutorStrategyContext::default()), + } } } diff --git a/crates/strategy/zksync/src/lib.rs b/crates/strategy/zksync/src/lib.rs index 8018e89be..344979a45 100644 --- a/crates/strategy/zksync/src/lib.rs +++ b/crates/strategy/zksync/src/lib.rs @@ -9,6 +9,8 @@ mod backend; mod cheatcode; mod executor; -pub use backend::ZksyncBackendStrategy; -pub use cheatcode::ZksyncCheatcodeInspectorStrategy; -pub use executor::{get_zksync_transaction_metadata, new_zkysnc_strategy, ZksyncExecutorStrategy}; +pub use backend::ZksyncBackendStrategyRunner; +pub use cheatcode::ZksyncCheatcodeInspectorStrategyRunner; +pub use executor::{ + get_zksync_transaction_metadata, ZksyncExecutorStrategyBuilder, ZksyncExecutorStrategyRunner, +}; diff --git a/crates/verify/src/utils.rs b/crates/verify/src/utils.rs index 60ba29715..cf83c0196 100644 --- a/crates/verify/src/utils.rs +++ b/crates/verify/src/utils.rs @@ -14,7 +14,7 @@ use foundry_compilers::artifacts::{BytecodeHash, CompactContractBytecode, EvmVer use foundry_config::Config; use foundry_evm::{ constants::DEFAULT_CREATE2_DEPLOYER, - executors::{strategy::Strategy, TracingExecutor}, + executors::{strategy::ExecutorStrategy, TracingExecutor}, opts::EvmOpts, }; use reqwest::Url; @@ -325,7 +325,7 @@ pub async fn get_tracing_executor( fork_blk_num: u64, evm_version: EvmVersion, evm_opts: EvmOpts, - strategy: Strategy, + strategy: ExecutorStrategy, ) -> Result<(Env, TracingExecutor)> { fork_config.fork_block_number = Some(fork_blk_num); fork_config.evm_version = evm_version; From e9781728218d3a27384d905b64f5c53a0c96a7ee Mon Sep 17 00:00:00 2001 From: Nisheeth Barthwal Date: Sun, 22 Dec 2024 20:39:06 +0100 Subject: [PATCH 5/6] clippy --- crates/evm/evm/src/executors/strategy.rs | 2 +- crates/strategy/zksync/src/backend.rs | 5 ++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/crates/evm/evm/src/executors/strategy.rs b/crates/evm/evm/src/executors/strategy.rs index 2a88deb3b..e27f57ce0 100644 --- a/crates/evm/evm/src/executors/strategy.rs +++ b/crates/evm/evm/src/executors/strategy.rs @@ -29,7 +29,7 @@ pub trait ExecutorStrategyContext: Debug + Send + Sync + Any { impl ExecutorStrategyContext for () { fn new_cloned(&self) -> Box { - Box::new(*self) + Box::new(()) } fn as_any_ref(&self) -> &dyn Any { diff --git a/crates/strategy/zksync/src/backend.rs b/crates/strategy/zksync/src/backend.rs index 76a9286ea..aadc55ef7 100644 --- a/crates/strategy/zksync/src/backend.rs +++ b/crates/strategy/zksync/src/backend.rs @@ -12,9 +12,8 @@ use foundry_evm_core::backend::{ BackendInner, Fork, ForkDB, FoundryEvmInMemoryDB, }; use foundry_zksync_core::{ - convert::ConvertH160, PaymasterParams, ACCOUNT_CODE_STORAGE_ADDRESS, - IMMUTABLE_SIMULATOR_STORAGE_ADDRESS, KNOWN_CODES_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 revm::{db::CacheDB, primitives::HashSet, DatabaseRef, JournaledState}; use serde::{Deserialize, Serialize}; From 28c431829431ada4ab850929fe8c911b053f1e05 Mon Sep 17 00:00:00 2001 From: Nisheeth Barthwal Date: Sun, 22 Dec 2024 21:31:30 +0100 Subject: [PATCH 6/6] fix test --- crates/forge/tests/it/test_helpers.rs | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/crates/forge/tests/it/test_helpers.rs b/crates/forge/tests/it/test_helpers.rs index 3446b8ec1..5ed3ed85b 100644 --- a/crates/forge/tests/it/test_helpers.rs +++ b/crates/forge/tests/it/test_helpers.rs @@ -3,7 +3,7 @@ use alloy_chains::NamedChain; use alloy_primitives::U256; use forge::{ - executors::strategy::new_evm_strategy, revm::primitives::SpecId, MultiContractRunner, + executors::strategy::ExecutorStrategy, revm::primitives::SpecId, MultiContractRunner, MultiContractRunnerBuilder, TestOptions, TestOptionsBuilder, }; use foundry_cli::utils; @@ -384,7 +384,7 @@ impl ForgeTestData { None, opts.local_evm_env(), opts, - new_evm_strategy(), + ExecutorStrategy::new_evm(), ) .unwrap() } @@ -401,7 +401,14 @@ impl ForgeTestData { self.base_runner() .with_fork(fork) - .build(self.project.root(), self.output.clone(), None, env, opts, new_evm_strategy()) + .build( + self.project.root(), + self.output.clone(), + None, + env, + opts, + ExecutorStrategy::new_evm(), + ) .unwrap() } }