Skip to content

Commit

Permalink
feat: get rid of zksync_select_fork_vm method (#813)
Browse files Browse the repository at this point in the history 
* feat: get rid of zksync_select_fork_vm method

* Fix return value
  • Loading branch information
@popzxc
popzxc authored Jan 8, 2025
1 parent b667ae0 commit e85875e
Show file tree
Hide file tree
Showing 3 changed files with 113 additions and 44 deletions.
10 changes: 0 additions & 10 deletions crates/cheatcodes/src/evm/fork.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -125,11 +125,6 @@ impl Cheatcode for selectForkCall {
persist_caller(ccx);
check_broadcast(ccx.state)?;

ccx.state.strategy.runner.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())
Expand Down Expand Up @@ -285,11 +280,6 @@ fn create_select_fork(ccx: &mut CheatsCtxt, url_or_alias: &str, block: Option<u6

let fork = create_fork_request(ccx, url_or_alias, block)?;
let id = ccx.ecx.db.create_fork(fork)?;
ccx.state.strategy.runner.zksync_select_fork_vm(
ccx.state.strategy.context.as_mut(),
ccx.ecx,
id,
);
ccx.ecx.db.select_fork(id, &mut ccx.ecx.env, &mut ccx.ecx.journaled_state)?;
Ok(id.abi_encode())
}
Expand Down
9 changes: 0 additions & 9 deletions crates/cheatcodes/src/strategy.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ use std::{any::Any, fmt::Debug, sync::Arc};

use alloy_primitives::{Address, TxKind};
use alloy_rpc_types::{TransactionInput, TransactionRequest};
use foundry_evm_core::backend::LocalForkId;
use revm::{
interpreter::{CallInputs, CallOutcome, CreateOutcome, Interpreter},
primitives::SignedAuthorization,
Expand Down Expand Up @@ -169,14 +168,6 @@ pub trait CheatcodeInspectorStrategyExt {
) -> Option<CallOutcome> {
None
}

fn zksync_select_fork_vm(
&self,
_ctx: &mut dyn CheatcodeInspectorStrategyContext,
_data: InnerEcx,
_fork_id: LocalForkId,
) {
}
}

#[derive(Debug, Default, Clone)]
Expand Down
138 changes: 113 additions & 25 deletions crates/strategy/zksync/src/cheatcode.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,9 +18,11 @@ use foundry_cheatcodes::{
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, zkVmCall, zkVmSkipCall,
self, createFork_0Call, createFork_1Call, createFork_2Call, createSelectFork_0Call,
createSelectFork_1Call, createSelectFork_2Call, dealCall, etchCall, getCodeCall,
getNonce_0Call, mockCallRevert_0Call, mockCall_0Call, resetNonceCall, rollCall,
selectForkCall, setNonceCall, setNonceUnsafeCall, warpCall, zkRegisterContractCall,
zkUseFactoryDepCall, zkUsePaymasterCall, zkVmCall, zkVmSkipCall,
},
};
use foundry_common::TransactionMaybeSigned;
Expand Down Expand Up @@ -271,31 +273,31 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
TypeId::of::<T>() == t
}

let ctx = get_context(ccx.state.strategy.context.as_mut());
let using_zk_vm = get_context(ccx.state.strategy.context.as_mut()).using_zk_vm;

// Try to downcast the cheatcode to a type that requires special handling.
// Note that some cheatcodes are only handled in zkEVM context.
// If no handler fires, we use the default execution logic.
match cheatcode.as_any().type_id() {
t if ctx.using_zk_vm && is::<etchCall>(t) => {
t if using_zk_vm && is::<etchCall>(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::<rollCall>(t) => {
t if using_zk_vm && is::<rollCall>(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::<warpCall>(t) => {
t if using_zk_vm && is::<warpCall>(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::<dealCall>(t) => {
t if using_zk_vm && is::<dealCall>(t) => {
let &dealCall { account, newBalance } = cheatcode.as_any().downcast_ref().unwrap();

let old_balance =
Expand All @@ -304,12 +306,12 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
ccx.state.eth_deals.push(record);
Ok(Default::default())
}
t if ctx.using_zk_vm && is::<resetNonceCall>(t) => {
t if using_zk_vm && is::<resetNonceCall>(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::<setNonceCall>(t) => {
t if using_zk_vm && is::<setNonceCall>(t) => {
let &setNonceCall { account, newNonce } =
cheatcode.as_any().downcast_ref().unwrap();

Expand All @@ -325,19 +327,19 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
foundry_zksync_core::cheatcodes::set_nonce(account, U256::from(newNonce), ccx.ecx);
Ok(Default::default())
}
t if ctx.using_zk_vm && is::<setNonceUnsafeCall>(t) => {
t if using_zk_vm && is::<setNonceUnsafeCall>(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::<getNonce_0Call>(t) => {
t if using_zk_vm && is::<getNonce_0Call>(t) => {
let &getNonce_0Call { account } = cheatcode.as_any().downcast_ref().unwrap();

let nonce = foundry_zksync_core::cheatcodes::get_nonce(account, ccx.ecx);
Ok(nonce.abi_encode())
}
t if ctx.using_zk_vm && is::<mockCall_0Call>(t) => {
t if using_zk_vm && is::<mockCall_0Call>(t) => {
let mockCall_0Call { callee, data, returnData } =
cheatcode.as_any().downcast_ref().unwrap();

Expand All @@ -353,7 +355,7 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
);
Ok(Default::default())
}
t if ctx.using_zk_vm && is::<mockCallRevert_0Call>(t) => {
t if using_zk_vm && is::<mockCallRevert_0Call>(t) => {
let mockCallRevert_0Call { callee, data, revertData } =
cheatcode.as_any().downcast_ref().unwrap();

Expand All @@ -374,6 +376,7 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
// We don't need to check for `using_zk_vm` since we pass it to `get_artifact_code`.
let getCodeCall { artifactPath } = cheatcode.as_any().downcast_ref().unwrap();

let ctx = get_context(ccx.state.strategy.context.as_mut());
Ok(get_artifact_code(
&ctx.dual_compiled_contracts,
ctx.using_zk_vm,
Expand All @@ -385,6 +388,7 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
}
t if is::<zkVmCall>(t) => {
let zkVmCall { enable } = cheatcode.as_any().downcast_ref().unwrap();
let ctx = get_context(ccx.state.strategy.context.as_mut());
if *enable {
self.select_zk_vm(ctx, ccx.ecx, None)
} else {
Expand All @@ -394,13 +398,14 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
}
t if is::<zkVmSkipCall>(t) => {
let zkVmSkipCall { .. } = cheatcode.as_any().downcast_ref().unwrap();
let ctx = get_context(ctx);
let ctx = get_context(ccx.state.strategy.context.as_mut());
ctx.skip_zk_vm = true;
Ok(Default::default())
}
t if is::<zkUsePaymasterCall>(t) => {
let zkUsePaymasterCall { paymaster_address, paymaster_input } =
cheatcode.as_any().downcast_ref().unwrap();
let ctx = get_context(ccx.state.strategy.context.as_mut());
ctx.paymaster_params = Some(ZkPaymasterData {
address: *paymaster_address,
input: paymaster_input.clone(),
Expand All @@ -410,6 +415,7 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
t if is::<zkUseFactoryDepCall>(t) => {
let zkUseFactoryDepCall { name } = cheatcode.as_any().downcast_ref().unwrap();
info!("Adding factory dependency: {:?}", name);
let ctx = get_context(ccx.state.strategy.context.as_mut());
ctx.zk_use_factory_deps.push(name.clone());
Ok(Default::default())
}
Expand All @@ -422,6 +428,7 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
zkBytecodeHash,
zkDeployedBytecode,
} = cheatcode.as_any().downcast_ref().unwrap();
let ctx = get_context(ccx.state.strategy.context.as_mut());

let zk_factory_deps = vec![]; //TODO: add argument to cheatcode
let new_contract = DualCompiledContract {
Expand Down Expand Up @@ -449,6 +456,97 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner

Ok(Default::default())
}
t if is::<selectForkCall>(t) => {
let selectForkCall { forkId } = cheatcode.as_any().downcast_ref().unwrap();
let ctx = get_context(ccx.state.strategy.context.as_mut());

// Re-implementation of `persist_caller` from `fork.rs`.
ccx.ecx.db.add_persistent_account(ccx.caller);

// Prepare storage.
self.select_fork_vm(ctx, ccx.ecx, *forkId);

// Apply cheatcode as usual.
cheatcode.dyn_apply(ccx, executor)
}
t if is::<createSelectFork_0Call>(t) => {
let createSelectFork_0Call { urlOrAlias } =
cheatcode.as_any().downcast_ref().unwrap();

// Re-implementation of `persist_caller` from `fork.rs`.
ccx.ecx.db.add_persistent_account(ccx.caller);

// Create fork.
let create_fork_cheatcode = createFork_0Call { urlOrAlias: urlOrAlias.clone() };

let encoded_fork_id = create_fork_cheatcode.dyn_apply(ccx, executor)?;
let fork_id = LocalForkId::abi_decode(&encoded_fork_id, true)?;

// Prepare storage.
{
let ctx = get_context(ccx.state.strategy.context.as_mut());
self.select_fork_vm(ctx, ccx.ecx, fork_id);
}

// Select fork
let select_fork_cheatcode = selectForkCall { forkId: fork_id };
select_fork_cheatcode.dyn_apply(ccx, executor)?;

// We need to return the fork ID.
Ok(encoded_fork_id)
}
t if is::<createSelectFork_1Call>(t) => {
let createSelectFork_1Call { urlOrAlias, blockNumber } =
cheatcode.as_any().downcast_ref().unwrap();

// Re-implementation of `persist_caller` from `fork.rs`.
ccx.ecx.db.add_persistent_account(ccx.caller);

// Create fork.
let create_fork_cheatcode =
createFork_1Call { urlOrAlias: urlOrAlias.clone(), blockNumber: *blockNumber };
let encoded_fork_id = create_fork_cheatcode.dyn_apply(ccx, executor)?;
let fork_id = LocalForkId::abi_decode(&encoded_fork_id, true)?;

// Prepare storage.
{
let ctx = get_context(ccx.state.strategy.context.as_mut());
self.select_fork_vm(ctx, ccx.ecx, fork_id);
}

// Select fork
let select_fork_cheatcode = selectForkCall { forkId: fork_id };
select_fork_cheatcode.dyn_apply(ccx, executor)?;

// We need to return the fork ID.
Ok(encoded_fork_id)
}
t if is::<createSelectFork_2Call>(t) => {
let createSelectFork_2Call { urlOrAlias, txHash } =
cheatcode.as_any().downcast_ref().unwrap();

// Re-implementation of `persist_caller` from `fork.rs`.
ccx.ecx.db.add_persistent_account(ccx.caller);

// Create fork.
let create_fork_cheatcode =
createFork_2Call { urlOrAlias: urlOrAlias.clone(), txHash: *txHash };
let encoded_fork_id = create_fork_cheatcode.dyn_apply(ccx, executor)?;
let fork_id = LocalForkId::abi_decode(&encoded_fork_id, true)?;

// Prepare storage.
{
let ctx = get_context(ccx.state.strategy.context.as_mut());
self.select_fork_vm(ctx, ccx.ecx, fork_id);
}

// Select fork
let select_fork_cheatcode = selectForkCall { forkId: fork_id };
select_fork_cheatcode.dyn_apply(ccx, executor)?;

// We need to return the fork ID.
Ok(encoded_fork_id)
}
_ => {
// Not custom, just invoke the default behavior
cheatcode.dyn_apply(ccx, executor)
Expand Down Expand Up @@ -1184,16 +1282,6 @@ impl CheatcodeInspectorStrategyExt for ZksyncCheatcodeInspectorStrategyRunner {
}
}
}

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);
}
}

impl ZksyncCheatcodeInspectorStrategyRunner {
Expand Down

0 comments on commit e85875e

Please sign in to comment.