From 40b9d9340e37e71836089b8556756da8d7218630 Mon Sep 17 00:00:00 2001 From: Igor Aleksanov Date: Mon, 6 Jan 2025 11:27:51 +0400 Subject: [PATCH] refactor: Move custom cheatcode logic into a generic handler --- crates/cheatcodes/src/evm.rs | 50 ++- crates/cheatcodes/src/evm/mock.rs | 9 +- crates/cheatcodes/src/fs.rs | 2 +- crates/cheatcodes/src/inspector.rs | 3 +- crates/cheatcodes/src/lib.rs | 9 +- crates/cheatcodes/src/strategy.rs | 182 +--------- crates/cheatcodes/src/test.rs | 63 +--- crates/strategy/zksync/src/cheatcode.rs | 455 +++++++++--------------- 8 files changed, 262 insertions(+), 511 deletions(-) diff --git a/crates/cheatcodes/src/evm.rs b/crates/cheatcodes/src/evm.rs index e5d6d3db5..e4a17fdd0 100644 --- a/crates/cheatcodes/src/evm.rs +++ b/crates/cheatcodes/src/evm.rs @@ -7,7 +7,7 @@ use crate::{ }; use alloy_consensus::TxEnvelope; use alloy_genesis::{Genesis, GenesisAccount}; -use alloy_primitives::{Address, B256, U256}; +use alloy_primitives::{Address, Bytes, B256, U256}; use alloy_rlp::Decodable; use alloy_sol_types::SolValue; use foundry_common::fs::{read_json_file, write_json_file}; @@ -18,7 +18,7 @@ use foundry_evm_core::{ }; use foundry_evm_traces::StackSnapshotType; use rand::Rng; -use revm::primitives::{Account, SpecId}; +use revm::primitives::{Account, Bytecode, SpecId, KECCAK_EMPTY}; use std::{collections::BTreeMap, path::Path}; mod record_debug_step; use record_debug_step::{convert_call_trace_to_debug_step, flatten_call_trace}; @@ -63,8 +63,7 @@ impl Cheatcode for addrCall { impl Cheatcode for getNonce_0Call { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { account } = self; - - ccx.state.strategy.runner.clone().cheatcode_get_nonce(ccx, *account) + get_nonce(ccx, account) } } @@ -350,7 +349,8 @@ impl Cheatcode for getBlobhashesCall { impl Cheatcode for rollCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { newHeight } = self; - ccx.state.strategy.runner.clone().cheatcode_roll(ccx, *newHeight) + ccx.ecx.env.block.number = *newHeight; + Ok(Default::default()) } } @@ -372,7 +372,8 @@ impl Cheatcode for txGasPriceCall { impl Cheatcode for warpCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { newTimestamp } = self; - ccx.state.strategy.runner.clone().cheatcode_warp(ccx, *newTimestamp) + ccx.ecx.env.block.timestamp = *newTimestamp; + Ok(Default::default()) } } @@ -407,7 +408,11 @@ impl Cheatcode for dealCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { account: address, newBalance: new_balance } = *self; - ccx.state.strategy.runner.clone().cheatcode_deal(ccx, address, new_balance) + 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 }; + ccx.state.eth_deals.push(record); + Ok(Default::default()) } } @@ -415,14 +420,26 @@ impl Cheatcode for etchCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { target, newRuntimeBytecode } = self; - ccx.state.strategy.runner.clone().cheatcode_etch(ccx, *target, newRuntimeBytecode) + ensure_not_precompile!(&target, ccx); + ccx.ecx.load_account(*target)?; + let bytecode = Bytecode::new_raw(Bytes::copy_from_slice(newRuntimeBytecode)); + ccx.ecx.journaled_state.set_code(*target, bytecode); + Ok(Default::default()) } } impl Cheatcode for resetNonceCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { account } = self; - ccx.state.strategy.runner.clone().cheatcode_reset_nonce(ccx, *account) + 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 + // to avoid hitting the case where account's code is None. + let empty = account.info.code_hash == KECCAK_EMPTY; + let nonce = if empty { 0 } else { 1 }; + account.info.nonce = nonce; + debug!(target: "cheatcodes", nonce, "reset"); + Ok(Default::default()) } } @@ -430,7 +447,16 @@ impl Cheatcode for setNonceCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { account, newNonce } = *self; - ccx.state.strategy.runner.clone().cheatcode_set_nonce(ccx, account, newNonce) + let account = journaled_account(ccx.ecx, account)?; + // nonce must increment only + let current = account.info.nonce; + ensure!( + newNonce >= current, + "new nonce ({newNonce}) must be strictly equal to or higher than the \ + account's current nonce ({current})" + ); + account.info.nonce = newNonce; + Ok(Default::default()) } } @@ -438,7 +464,9 @@ impl Cheatcode for setNonceUnsafeCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { account, newNonce } = *self; - ccx.state.strategy.runner.clone().cheatcode_set_nonce_unsafe(ccx, account, newNonce) + let account = journaled_account(ccx.ecx, account)?; + account.info.nonce = newNonce; + Ok(Default::default()) } } diff --git a/crates/cheatcodes/src/evm/mock.rs b/crates/cheatcodes/src/evm/mock.rs index e720e0b43..7b716bbeb 100644 --- a/crates/cheatcodes/src/evm/mock.rs +++ b/crates/cheatcodes/src/evm/mock.rs @@ -15,7 +15,9 @@ impl Cheatcode for clearMockedCallsCall { impl Cheatcode for mockCall_0Call { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { let Self { callee, data, returnData } = self; - ccx.state.strategy.runner.clone().cheatcode_mock_call(ccx, *callee, data, returnData) + let _ = make_acc_non_empty(callee, ccx.ecx)?; + mock_call(ccx.state, callee, data, None, returnData, InstructionResult::Return); + Ok(Default::default()) } } @@ -83,7 +85,10 @@ 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.runner.clone().cheatcode_mock_call_revert(ccx, *callee, data, revertData) + + let _ = make_acc_non_empty(callee, ccx.ecx)?; + mock_call(ccx.state, callee, data, None, revertData, InstructionResult::Revert); + Ok(Default::default()) } } diff --git a/crates/cheatcodes/src/fs.rs b/crates/cheatcodes/src/fs.rs index 96e53704c..14a3745c0 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.runner.get_artifact_code(state, path, false) + Ok(crate::fs::get_artifact_code(state, path, false)?.abi_encode()) } } diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index ad1f757e2..fd79f8831 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -2256,7 +2256,8 @@ fn apply_dispatch( } // Apply the cheatcode. - let mut result = cheat.dyn_apply(ccx, executor); + let runner = ccx.state.strategy.runner.new_cloned(); + let mut result = runner.apply_full(cheat, ccx, executor); // Format the error message to include the cheatcode name. if let Err(e) = &mut result { diff --git a/crates/cheatcodes/src/lib.rs b/crates/cheatcodes/src/lib.rs index 33dec1c33..3375ff057 100644 --- a/crates/cheatcodes/src/lib.rs +++ b/crates/cheatcodes/src/lib.rs @@ -96,12 +96,14 @@ pub(crate) trait Cheatcode: CheatcodeDef + DynCheatcode { } } -pub(crate) trait DynCheatcode: 'static { +pub trait DynCheatcode: 'static + std::any::Any { fn cheatcode(&self) -> &'static spec::Cheatcode<'static>; fn as_debug(&self) -> &dyn std::fmt::Debug; fn dyn_apply(&self, ccx: &mut CheatsCtxt, executor: &mut dyn CheatcodesExecutor) -> Result; + + fn as_any(&self) -> &dyn std::any::Any; } impl DynCheatcode for T { @@ -119,6 +121,11 @@ impl DynCheatcode for T { fn dyn_apply(&self, ccx: &mut CheatsCtxt, executor: &mut dyn CheatcodesExecutor) -> Result { self.apply_full(ccx, executor) } + + #[inline] + fn as_any(&self) -> &dyn std::any::Any { + self + } } impl dyn DynCheatcode { diff --git a/crates/cheatcodes/src/strategy.rs b/crates/cheatcodes/src/strategy.rs index 42d1f6fbf..aa0d897bc 100644 --- a/crates/cheatcodes/src/strategy.rs +++ b/crates/cheatcodes/src/strategy.rs @@ -1,21 +1,18 @@ use std::{any::Any, fmt::Debug, sync::Arc}; -use alloy_primitives::{Address, Bytes, FixedBytes, TxKind, U256}; +use alloy_primitives::{Address, TxKind}; use alloy_rpc_types::{TransactionInput, TransactionRequest}; -use alloy_sol_types::SolValue; use foundry_evm_core::backend::LocalForkId; use revm::{ - interpreter::{CallInputs, CallOutcome, CreateOutcome, InstructionResult, Interpreter}, - primitives::{Bytecode, SignedAuthorization, KECCAK_EMPTY}, + interpreter::{CallInputs, CallOutcome, CreateOutcome, Interpreter}, + primitives::SignedAuthorization, }; use crate::{ - evm::{self, journaled_account, mock::make_acc_non_empty, DealRecord}, inspector::{check_if_fixed_gas_limit, CommonCreateInput, Ecx, InnerEcx}, - mock_call, script::Broadcast, BroadcastableTransaction, BroadcastableTransactions, Cheatcodes, CheatcodesExecutor, - CheatsConfig, CheatsCtxt, Result, + CheatsConfig, CheatsCtxt, DynCheatcode, Result, }; /// Represents the context for [CheatcodeInspectorStrategy]. @@ -73,129 +70,17 @@ pub trait CheatcodeInspectorStrategyRunner: fn new_cloned(&self) -> Box; - /// Get nonce. - 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(&self, _ctx: &mut dyn CheatcodeInspectorStrategyContext) {} - - /// Cheatcode: roll. - 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(&self, ccx: &mut CheatsCtxt, new_timestamp: U256) -> Result { - ccx.ecx.env.block.timestamp = new_timestamp; - Ok(Default::default()) - } - - /// Cheatcode: deal. - 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 }; - ccx.state.eth_deals.push(record); - Ok(Default::default()) - } - - /// Cheatcode: etch. - fn cheatcode_etch( + fn apply_full( &self, + cheatcode: &dyn DynCheatcode, ccx: &mut CheatsCtxt, - target: Address, - new_runtime_bytecode: &Bytes, + executor: &mut dyn CheatcodesExecutor, ) -> Result { - ensure_not_precompile!(&target, ccx); - ccx.ecx.load_account(target)?; - let bytecode = Bytecode::new_raw(Bytes::copy_from_slice(new_runtime_bytecode)); - ccx.ecx.journaled_state.set_code(target, bytecode); - Ok(Default::default()) - } - - /// Cheatcode: getNonce. - fn cheatcode_get_nonce(&self, ccx: &mut CheatsCtxt, address: Address) -> Result { - evm::get_nonce(ccx, &address) + cheatcode.dyn_apply(ccx, executor) } - /// Cheatcode: resetNonce. - 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 - // to avoid hitting the case where account's code is None. - let empty = account.info.code_hash == KECCAK_EMPTY; - let nonce = if empty { 0 } else { 1 }; - account.info.nonce = nonce; - debug!(target: "cheatcodes", nonce, "reset"); - Ok(Default::default()) - } - - /// Cheatcode: setNonce. - fn cheatcode_set_nonce( - &self, - ccx: &mut CheatsCtxt, - account: Address, - new_nonce: u64, - ) -> Result { - let account = journaled_account(ccx.ecx, account)?; - // nonce must increment only - let current = account.info.nonce; - ensure!( - new_nonce >= current, - "new nonce ({new_nonce}) must be strictly equal to or higher than the \ - account's current nonce ({current})" - ); - account.info.nonce = new_nonce; - Ok(Default::default()) - } - - /// Cheatcode: setNonceUnsafe. - fn cheatcode_set_nonce_unsafe( - &self, - ccx: &mut CheatsCtxt, - account: Address, - new_nonce: u64, - ) -> Result { - let account = journaled_account(ccx.ecx, account)?; - account.info.nonce = new_nonce; - Ok(Default::default()) - } - - /// Mocks a call to return with a value. - fn cheatcode_mock_call( - &self, - ccx: &mut CheatsCtxt, - callee: Address, - data: &Bytes, - return_data: &Bytes, - ) -> Result { - let _ = make_acc_non_empty(&callee, ccx.ecx)?; - mock_call(ccx.state, &callee, data, None, return_data, InstructionResult::Return); - Ok(Default::default()) - } - - /// Mocks a call to revert with a value. - fn cheatcode_mock_call_revert( - &self, - ccx: &mut CheatsCtxt, - callee: Address, - data: &Bytes, - revert_data: &Bytes, - ) -> Result { - let _ = make_acc_non_empty(&callee, ccx.ecx)?; - mock_call(ccx.state, &callee, data, None, revert_data, InstructionResult::Revert); - Ok(Default::default()) - } - - /// Retrieve artifact code. - fn get_artifact_code(&self, state: &Cheatcodes, path: &str, deployed: bool) -> Result { - Ok(crate::fs::get_artifact_code(state, path, deployed)?.abi_encode()) - } + /// Called when the main test or script contract is deployed. + fn base_contract_deployed(&self, _ctx: &mut dyn CheatcodeInspectorStrategyContext) {} /// Record broadcastable transaction during CREATE. fn record_broadcastable_create_transactions( @@ -242,53 +127,6 @@ pub trait CheatcodeInspectorStrategyRunner: /// We define this in our fork pub trait CheatcodeInspectorStrategyExt { - fn zksync_cheatcode_skip_zkvm( - &self, - _ctx: &mut dyn CheatcodeInspectorStrategyContext, - ) -> Result { - Ok(Default::default()) - } - - fn zksync_cheatcode_set_paymaster( - &self, - _ctx: &mut dyn CheatcodeInspectorStrategyContext, - _paymaster_address: Address, - _paymaster_input: &Bytes, - ) -> Result { - Ok(Default::default()) - } - - 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( - &self, - _ctx: &mut dyn CheatcodeInspectorStrategyContext, - _name: String, - _zk_bytecode_hash: FixedBytes<32>, - _zk_deployed_bytecode: Vec, - _zk_factory_deps: Vec>, - _evm_bytecode_hash: FixedBytes<32>, - _evm_deployed_bytecode: Vec, - _evm_bytecode: Vec, - ) -> Result { - Ok(Default::default()) - } - - fn zksync_cheatcode_select_zk_vm( - &self, - _ctx: &mut dyn CheatcodeInspectorStrategyContext, - _data: InnerEcx, - _enable: bool, - ) { - } - fn zksync_record_create_address( &self, _ctx: &mut dyn CheatcodeInspectorStrategyContext, diff --git a/crates/cheatcodes/src/test.rs b/crates/cheatcodes/src/test.rs index ae59439fc..ad62f1410 100644 --- a/crates/cheatcodes/src/test.rs +++ b/crates/cheatcodes/src/test.rs @@ -12,67 +12,42 @@ pub(crate) mod assume; pub(crate) mod expect; impl Cheatcode for zkVmCall { - fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { - let Self { enable } = *self; - - ccx.state.strategy.runner.zksync_cheatcode_select_zk_vm( - ccx.state.strategy.context.as_mut(), - ccx.ecx, - enable, - ); - + fn apply_stateful(&self, _ccx: &mut CheatsCtxt) -> Result { + // Does nothing by default. + // ZK-related logic is implemented in the corresponding strategy object. Ok(Default::default()) } } impl Cheatcode for zkVmSkipCall { - fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { - ccx.state.strategy.runner.zksync_cheatcode_skip_zkvm(ccx.state.strategy.context.as_mut()) + fn apply_stateful(&self, _ccx: &mut CheatsCtxt) -> Result { + // Does nothing by default. + // ZK-related logic is implemented in the corresponding strategy object. + Ok(Default::default()) } } impl Cheatcode for zkUsePaymasterCall { - fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { - let Self { paymaster_address, paymaster_input } = self; - ccx.state.strategy.runner.zksync_cheatcode_set_paymaster( - ccx.state.strategy.context.as_mut(), - *paymaster_address, - paymaster_input, - ) + fn apply_stateful(&self, _ccx: &mut CheatsCtxt) -> Result { + // Does nothing by default. + // ZK-related logic is implemented in the corresponding strategy object. + Ok(Default::default()) } } impl Cheatcode for zkUseFactoryDepCall { - fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { - let Self { name } = self; - ccx.state - .strategy - .runner - .zksync_cheatcode_use_factory_deps(ccx.state.strategy.context.as_mut(), name.clone()) + fn apply_stateful(&self, _ccx: &mut CheatsCtxt) -> Result { + // Does nothing by default. + // ZK-related logic is implemented in the corresponding strategy object. + Ok(Default::default()) } } impl Cheatcode for zkRegisterContractCall { - fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { - let Self { - name, - evmBytecodeHash, - evmDeployedBytecode, - evmBytecode, - zkBytecodeHash, - zkDeployedBytecode, - } = self; - - ccx.state.strategy.runner.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(), - ) + fn apply_stateful(&self, _ccx: &mut CheatsCtxt) -> Result { + // Does nothing by default. + // ZK-related logic is implemented in the corresponding strategy object. + Ok(Default::default()) } } diff --git a/crates/strategy/zksync/src/cheatcode.rs b/crates/strategy/zksync/src/cheatcode.rs index dc7ca43c0..7163ef23b 100644 --- a/crates/strategy/zksync/src/cheatcode.rs +++ b/crates/strategy/zksync/src/cheatcode.rs @@ -1,7 +1,7 @@ -use std::{fs, path::PathBuf, sync::Arc}; +use std::{any::TypeId, fs, path::PathBuf, sync::Arc}; use alloy_json_abi::ContractObject; -use alloy_primitives::{keccak256, map::HashMap, Address, Bytes, FixedBytes, TxKind, B256, U256}; +use alloy_primitives::{keccak256, map::HashMap, Address, Bytes, TxKind, B256, U256}; use alloy_rpc_types::{ request::{TransactionInput, TransactionRequest}, serde_helpers::WithOtherFields, @@ -15,7 +15,13 @@ use foundry_cheatcodes::{ EvmCheatcodeInspectorStrategyRunner, }, Broadcast, BroadcastableTransaction, BroadcastableTransactions, Cheatcodes, CheatcodesExecutor, - CheatsConfig, CheatsCtxt, CommonCreateInput, DealRecord, Ecx, Error, InnerEcx, Result, Vm, + CheatsConfig, CheatsCtxt, CommonCreateInput, DealRecord, DynCheatcode, Ecx, Error, InnerEcx, + Result, + Vm::{ + self, dealCall, etchCall, getCodeCall, getNonce_0Call, mockCallRevert_0Call, + mockCall_0Call, resetNonceCall, rollCall, setNonceCall, setNonceUnsafeCall, warpCall, + zkRegisterContractCall, zkUseFactoryDepCall, zkUsePaymasterCall, zkVmSkipCall, + }, }; use foundry_common::TransactionMaybeSigned; use foundry_config::fs_permissions::FsAccessKind; @@ -244,17 +250,6 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner Box::new(self.clone()) } - 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); - } - - let nonce = foundry_zksync_core::nonce(address, ccx.ecx) as u64; - Ok(nonce) - } - fn base_contract_deployed(&self, ctx: &mut dyn CheatcodeInspectorStrategyContext) { let ctx = get_context(ctx); @@ -264,200 +259,187 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner ctx.zk_persist_nonce_update.persist_next(); } - fn cheatcode_get_nonce( - &self, - ccx: &mut CheatsCtxt<'_, '_, '_, '_>, - address: Address, - ) -> foundry_cheatcodes::Result { - 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()); - } - - let nonce = foundry_zksync_core::cheatcodes::get_nonce(address, ccx.ecx); - Ok(nonce.abi_encode()) - } - - 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); - } - - ccx.ecx.env.block.number = new_height; - foundry_zksync_core::cheatcodes::roll(new_height, ccx.ecx); - Ok(Default::default()) - } - - 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); - } - - ccx.ecx.env.block.number = new_timestamp; - foundry_zksync_core::cheatcodes::warp(new_timestamp, ccx.ecx); - Ok(Default::default()) - } - - fn cheatcode_deal( + fn apply_full( &self, + cheatcode: &dyn DynCheatcode, ccx: &mut CheatsCtxt<'_, '_, '_, '_>, - address: Address, - new_balance: U256, - ) -> Result { - let ctx = get_context(ccx.state.strategy.context.as_mut()); - - if !ctx.using_zk_vm { - return self.evm.cheatcode_deal(ccx, address, new_balance); - } - - let old_balance = foundry_zksync_core::cheatcodes::deal(address, new_balance, ccx.ecx); - let record = DealRecord { address, old_balance, new_balance }; - ccx.state.eth_deals.push(record); - Ok(Default::default()) - } - - fn cheatcode_etch( - &self, - ccx: &mut CheatsCtxt<'_, '_, '_, '_>, - target: Address, - new_runtime_bytecode: &Bytes, - ) -> Result { - 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); - } - - foundry_zksync_core::cheatcodes::etch(target, new_runtime_bytecode, ccx.ecx); - Ok(Default::default()) - } - - fn cheatcode_reset_nonce( - &self, - ccx: &mut CheatsCtxt<'_, '_, '_, '_>, - account: Address, - ) -> Result { - let ctx = get_context(ccx.state.strategy.context.as_mut()); - - if !ctx.using_zk_vm { - return self.evm.cheatcode_reset_nonce(ccx, account); - } - - foundry_zksync_core::cheatcodes::set_nonce(account, U256::ZERO, ccx.ecx); - Ok(Default::default()) - } - - fn cheatcode_set_nonce( - &self, - ccx: &mut CheatsCtxt<'_, '_, '_, '_>, - account: Address, - new_nonce: u64, + executor: &mut dyn CheatcodesExecutor, ) -> Result { - 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); + fn is(t: TypeId) -> bool { + TypeId::of::() == t } - // nonce must increment only - let current = foundry_zksync_core::cheatcodes::get_nonce(account, ccx.ecx); - if U256::from(new_nonce) < current { - return Err(fmt_err!( - "new nonce ({new_nonce}) must be strictly equal to or higher than the \ - account's current nonce ({current})" - )); - } - - foundry_zksync_core::cheatcodes::set_nonce(account, U256::from(new_nonce), ccx.ecx); - Ok(Default::default()) - } - - fn cheatcode_set_nonce_unsafe( - &self, - ccx: &mut CheatsCtxt<'_, '_, '_, '_>, - account: Address, - new_nonce: u64, - ) -> Result { 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); - } + match cheatcode.as_any().type_id() { + t if ctx.using_zk_vm && is::(t) => { + let etchCall { target, newRuntimeBytecode } = + cheatcode.as_any().downcast_ref().unwrap(); + foundry_zksync_core::cheatcodes::etch(*target, newRuntimeBytecode, ccx.ecx); + Ok(Default::default()) + } + t if ctx.using_zk_vm && is::(t) => { + let &rollCall { newHeight } = cheatcode.as_any().downcast_ref().unwrap(); + ccx.ecx.env.block.number = newHeight; + foundry_zksync_core::cheatcodes::roll(newHeight, ccx.ecx); + Ok(Default::default()) + } + t if ctx.using_zk_vm && is::(t) => { + let &warpCall { newTimestamp } = cheatcode.as_any().downcast_ref().unwrap(); + ccx.ecx.env.block.number = newTimestamp; + foundry_zksync_core::cheatcodes::warp(newTimestamp, ccx.ecx); + Ok(Default::default()) + } + t if ctx.using_zk_vm && is::(t) => { + let &dealCall { account, newBalance } = cheatcode.as_any().downcast_ref().unwrap(); + + let old_balance = + foundry_zksync_core::cheatcodes::deal(account, newBalance, ccx.ecx); + let record = DealRecord { address: account, old_balance, new_balance: newBalance }; + ccx.state.eth_deals.push(record); + Ok(Default::default()) + } + t if ctx.using_zk_vm && is::(t) => { + let &resetNonceCall { account } = cheatcode.as_any().downcast_ref().unwrap(); + foundry_zksync_core::cheatcodes::set_nonce(account, U256::ZERO, ccx.ecx); + Ok(Default::default()) + } + t if ctx.using_zk_vm && is::(t) => { + let &setNonceCall { account, newNonce } = + cheatcode.as_any().downcast_ref().unwrap(); + + // nonce must increment only + let current = foundry_zksync_core::cheatcodes::get_nonce(account, ccx.ecx); + if U256::from(newNonce) < current { + return Err(fmt_err!( + "new nonce ({newNonce}) must be strictly equal to or higher than the \ + account's current nonce ({current})" + )); + } - foundry_zksync_core::cheatcodes::set_nonce(account, U256::from(new_nonce), ccx.ecx); - Ok(Default::default()) - } + foundry_zksync_core::cheatcodes::set_nonce(account, U256::from(newNonce), ccx.ecx); + Ok(Default::default()) + } + t if ctx.using_zk_vm && is::(t) => { + let &setNonceUnsafeCall { account, newNonce } = + cheatcode.as_any().downcast_ref().unwrap(); + foundry_zksync_core::cheatcodes::set_nonce(account, U256::from(newNonce), ccx.ecx); + Ok(Default::default()) + } + t if ctx.using_zk_vm && is::(t) => { + let &getNonce_0Call { account } = cheatcode.as_any().downcast_ref().unwrap(); - fn cheatcode_mock_call( - &self, - ccx: &mut CheatsCtxt<'_, '_, '_, '_>, - callee: Address, - data: &Bytes, - return_data: &Bytes, - ) -> Result { - let ctx = get_context(ccx.state.strategy.context.as_mut()); + let nonce = foundry_zksync_core::cheatcodes::get_nonce(account, ccx.ecx); + Ok(nonce.abi_encode()) + } + t if ctx.using_zk_vm && is::(t) => { + let mockCall_0Call { callee, data, returnData } = + cheatcode.as_any().downcast_ref().unwrap(); + + let _ = foundry_cheatcodes::make_acc_non_empty(callee, ccx.ecx)?; + foundry_zksync_core::cheatcodes::set_mocked_account(*callee, ccx.ecx, ccx.caller); + foundry_cheatcodes::mock_call( + ccx.state, + callee, + data, + None, + returnData, + InstructionResult::Return, + ); + Ok(Default::default()) + } + t if ctx.using_zk_vm && is::(t) => { + let mockCallRevert_0Call { callee, data, revertData } = + cheatcode.as_any().downcast_ref().unwrap(); + + let _ = make_acc_non_empty(callee, ccx.ecx)?; + foundry_zksync_core::cheatcodes::set_mocked_account(*callee, ccx.ecx, ccx.caller); + // not calling + foundry_cheatcodes::mock_call( + ccx.state, + callee, + data, + None, + revertData, + InstructionResult::Revert, + ); + Ok(Default::default()) + } + t if ctx.using_zk_vm && is::(t) => { + let getCodeCall { artifactPath } = cheatcode.as_any().downcast_ref().unwrap(); - if !ctx.using_zk_vm { - return self.evm.cheatcode_mock_call(ccx, callee, data, return_data); - } + let ctx = get_context_ref(ccx.state.strategy.context.as_ref()); + Ok(get_artifact_code( + &ctx.dual_compiled_contracts, + ctx.using_zk_vm, + &ccx.state.config, + artifactPath, + false, + )? + .abi_encode()) + } + t if is::(t) => { + let zkVmSkipCall { .. } = cheatcode.as_any().downcast_ref().unwrap(); + let ctx = get_context(ctx); + ctx.skip_zk_vm = true; + Ok(Default::default()) + } + t if is::(t) => { + let zkUsePaymasterCall { paymaster_address, paymaster_input } = + cheatcode.as_any().downcast_ref().unwrap(); + ctx.paymaster_params = Some(ZkPaymasterData { + address: *paymaster_address, + input: paymaster_input.clone(), + }); + Ok(Default::default()) + } + t if is::(t) => { + let zkUseFactoryDepCall { name } = cheatcode.as_any().downcast_ref().unwrap(); + info!("Adding factory dependency: {:?}", name); + ctx.zk_use_factory_deps.push(name.clone()); + Ok(Default::default()) + } + t if is::(t) => { + let zkRegisterContractCall { + name, + evmBytecodeHash, + evmDeployedBytecode, + evmBytecode, + zkBytecodeHash, + zkDeployedBytecode, + } = cheatcode.as_any().downcast_ref().unwrap(); + + let zk_factory_deps = vec![]; //TODO: add argument to cheatcode + let new_contract = DualCompiledContract { + name: name.clone(), + zk_bytecode_hash: H256(zkBytecodeHash.0), + zk_deployed_bytecode: zkDeployedBytecode.to_vec(), + zk_factory_deps, + evm_bytecode_hash: *evmBytecodeHash, + evm_deployed_bytecode: evmDeployedBytecode.to_vec(), + evm_bytecode: evmBytecode.to_vec(), + }; - let _ = foundry_cheatcodes::make_acc_non_empty(&callee, ccx.ecx)?; - foundry_zksync_core::cheatcodes::set_mocked_account(callee, ccx.ecx, ccx.caller); - foundry_cheatcodes::mock_call( - ccx.state, - &callee, - data, - None, - return_data, - InstructionResult::Return, - ); - Ok(Default::default()) - } + 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 + }) { + warn!( + name = existing.name, + "contract already exists with the given bytecode hashes" + ); + return Ok(Default::default()) + } - fn cheatcode_mock_call_revert( - &self, - ccx: &mut CheatsCtxt<'_, '_, '_, '_>, - callee: Address, - data: &Bytes, - revert_data: &Bytes, - ) -> Result { - let ctx = get_context(ccx.state.strategy.context.as_mut()); + ctx.dual_compiled_contracts.push(new_contract); - if !ctx.using_zk_vm { - return self.evm.cheatcode_mock_call_revert(ccx, callee, data, revert_data); + Ok(Default::default()) + } + _ => { + // Not custom, just invoke the default behavior + cheatcode.dyn_apply(ccx, executor) + } } - - let _ = make_acc_non_empty(&callee, ccx.ecx)?; - foundry_zksync_core::cheatcodes::set_mocked_account(callee, ccx.ecx, ccx.caller); - // not calling - foundry_cheatcodes::mock_call( - ccx.state, - &callee, - data, - None, - revert_data, - InstructionResult::Revert, - ); - Ok(Default::default()) - } - - 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( - &ctx.dual_compiled_contracts, - ctx.using_zk_vm, - &state.config, - path, - deployed, - )? - .abi_encode()) } fn record_broadcastable_create_transactions( @@ -742,77 +724,6 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner } impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategyRunner { - 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( - &self, - ctx: &mut dyn CheatcodeInspectorStrategyContext, - paymaster_address: Address, - paymaster_input: &Bytes, - ) -> Result { - 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( - &self, - ctx: &mut dyn CheatcodeInspectorStrategyContext, - name: String, - ) -> foundry_cheatcodes::Result { - let ctx = get_context(ctx); - - info!("Adding factory dependency: {:?}", name); - ctx.zk_use_factory_deps.push(name); - Ok(Default::default()) - } - - fn zksync_cheatcode_register_contract( - &self, - ctx: &mut dyn CheatcodeInspectorStrategyContext, - name: String, - zk_bytecode_hash: FixedBytes<32>, - zk_deployed_bytecode: Vec, - zk_factory_deps: Vec>, - evm_bytecode_hash: FixedBytes<32>, - 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), - zk_deployed_bytecode, - zk_factory_deps, - evm_bytecode_hash, - evm_deployed_bytecode, - evm_bytecode, - }; - - 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 - }) { - warn!(name = existing.name, "contract already exists with the given bytecode hashes"); - return Ok(Default::default()) - } - - ctx.dual_compiled_contracts.push(new_contract); - - Ok(Default::default()) - } - fn zksync_record_create_address( &self, ctx: &mut dyn CheatcodeInspectorStrategyContext, @@ -1269,20 +1180,6 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategyRunner { let ctx = get_context(ctx); self.select_fork_vm(ctx, data, fork_id); } - - 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(ctx, data, None) - } else { - self.select_evm(ctx, data); - } - } } impl ZksyncCheatcodeInspectorStrategyRunner {