Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor: Simplify strategy traits #818

Merged
merged 1 commit into from
Jan 8, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 3 additions & 8 deletions crates/cheatcodes/src/inspector.rs
Original file line number Diff line number Diff line change
Expand Up @@ -801,9 +801,7 @@ impl Cheatcodes {
}]);
}

if let Some(result) =
self.strategy.runner.clone().zksync_try_create(self, ecx, &input, executor)
{
if let Some(result) = self.strategy.runner.zksync_try_create(self, ecx, &input, executor) {
return Some(result);
}

Expand Down Expand Up @@ -1222,9 +1220,7 @@ where {
}]);
}

if let Some(result) =
self.strategy.runner.clone().zksync_try_call(self, ecx, call, executor)
{
if let Some(result) = self.strategy.runner.zksync_try_call(self, ecx, call, executor) {
return Some(result);
}

Expand Down Expand Up @@ -2285,8 +2281,7 @@ fn apply_dispatch(
}

// Apply the cheatcode.
let runner = ccx.state.strategy.runner.new_cloned();
let mut result = runner.apply_full(cheat, ccx, executor);
let mut result = ccx.state.strategy.runner.apply_full(cheat, ccx, executor);

// Format the error message to include the cheatcode name.
if let Err(e) = &mut result {
Expand Down
29 changes: 4 additions & 25 deletions crates/cheatcodes/src/strategy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,33 +42,26 @@ impl CheatcodeInspectorStrategyContext for () {
#[derive(Debug)]
pub struct CheatcodeInspectorStrategy {
/// Strategy runner.
pub runner: Box<dyn CheatcodeInspectorStrategyRunner>,
pub runner: &'static dyn CheatcodeInspectorStrategyRunner,
/// Strategy context.
pub context: Box<dyn CheatcodeInspectorStrategyContext>,
}

impl CheatcodeInspectorStrategy {
pub fn new_evm() -> Self {
Self {
runner: Box::new(EvmCheatcodeInspectorStrategyRunner::default()),
context: Box::new(()),
}
Self { runner: &EvmCheatcodeInspectorStrategyRunner, context: Box::new(()) }
}
}

impl Clone for CheatcodeInspectorStrategy {
fn clone(&self) -> Self {
Self { runner: self.runner.new_cloned(), context: self.context.new_cloned() }
Self { runner: self.runner, context: self.context.new_cloned() }
}
}

pub trait CheatcodeInspectorStrategyRunner:
Debug + Send + Sync + CheatcodeInspectorStrategyExt
{
fn name(&self) -> &'static str;

fn new_cloned(&self) -> Box<dyn CheatcodeInspectorStrategyRunner>;

fn apply_full(
&self,
cheatcode: &dyn DynCheatcode,
Expand Down Expand Up @@ -171,17 +164,9 @@ pub trait CheatcodeInspectorStrategyExt {
}

#[derive(Debug, Default, Clone)]
pub struct EvmCheatcodeInspectorStrategyRunner {}
pub struct EvmCheatcodeInspectorStrategyRunner;

impl CheatcodeInspectorStrategyRunner for EvmCheatcodeInspectorStrategyRunner {
fn name(&self) -> &'static str {
"evm"
}

fn new_cloned(&self) -> Box<dyn CheatcodeInspectorStrategyRunner> {
Box::new(self.clone())
}

fn record_broadcastable_create_transactions(
&self,
_ctx: &mut dyn CheatcodeInspectorStrategyContext,
Expand Down Expand Up @@ -256,11 +241,5 @@ impl CheatcodeInspectorStrategyRunner for EvmCheatcodeInspectorStrategyRunner {

impl CheatcodeInspectorStrategyExt for EvmCheatcodeInspectorStrategyRunner {}

impl Clone for Box<dyn CheatcodeInspectorStrategyRunner> {
fn clone(&self) -> Self {
self.new_cloned()
}
}

struct _ObjectSafe0(dyn CheatcodeInspectorStrategyRunner);
struct _ObjectSafe1(dyn CheatcodeInspectorStrategyExt);
7 changes: 1 addition & 6 deletions crates/evm/core/src/backend/cow.rs
Original file line number Diff line number Diff line change
Expand Up @@ -72,12 +72,7 @@ impl<'a> CowBackend<'a> {
self.is_initialized = false;
self.spec_id = env.handler_cfg.spec_id;

self.backend.strategy.runner.clone().inspect(
self.backend.to_mut(),
env,
inspector,
inspect_ctx,
)
self.backend.strategy.runner.inspect(self.backend.to_mut(), env, inspector, inspect_ctx)
}

/// Returns whether there was a state snapshot failure in the backend.
Expand Down
2 changes: 1 addition & 1 deletion crates/evm/core/src/backend/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -794,7 +794,7 @@ impl Backend {
inspect_ctx: Box<dyn Any>,
) -> eyre::Result<ResultAndState> {
self.initialize(env);
self.strategy.runner.clone().inspect(self, env, inspector, inspect_ctx)
self.strategy.runner.inspect(self, env, inspector, inspect_ctx)
}

/// Returns the `EnvWithHandlerCfg` with the current `spec_id` set.
Expand Down
24 changes: 3 additions & 21 deletions crates/evm/core/src/backend/strategy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,29 +46,25 @@ impl BackendStrategyContext for () {
#[derive(Debug)]
pub struct BackendStrategy {
/// Strategy runner.
pub runner: Box<dyn BackendStrategyRunner>,
pub runner: &'static dyn BackendStrategyRunner,
/// Strategy context.
pub context: Box<dyn BackendStrategyContext>,
}

impl BackendStrategy {
/// Create a new instance of [BackendStrategy]
pub fn new_evm() -> Self {
Self { runner: Box::new(EvmBackendStrategyRunner), context: Box::new(()) }
Self { runner: &EvmBackendStrategyRunner, context: Box::new(()) }
}
}

impl Clone for BackendStrategy {
fn clone(&self) -> Self {
Self { runner: self.runner.new_cloned(), context: self.context.new_cloned() }
Self { runner: self.runner, context: self.context.new_cloned() }
}
}

pub trait BackendStrategyRunner: Debug + Send + Sync + BackendStrategyRunnerExt {
fn name(&self) -> &'static str;

fn new_cloned(&self) -> Box<dyn BackendStrategyRunner>;

fn inspect(
&self,
backend: &mut Backend,
Expand Down Expand Up @@ -128,14 +124,6 @@ struct _ObjectSafe(dyn BackendStrategyRunner);
pub struct EvmBackendStrategyRunner;

impl BackendStrategyRunner for EvmBackendStrategyRunner {
fn name(&self) -> &'static str {
"evm"
}

fn new_cloned(&self) -> Box<dyn BackendStrategyRunner> {
Box::new(self.clone())
}

fn inspect(
&self,
backend: &mut Backend,
Expand Down Expand Up @@ -298,9 +286,3 @@ impl EvmBackendMergeStrategy {
fork_db.accounts.insert(addr, acc);
}
}

impl Clone for Box<dyn BackendStrategyRunner> {
fn clone(&self) -> Self {
self.new_cloned()
}
}
4 changes: 2 additions & 2 deletions crates/evm/evm/src/executors/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -248,7 +248,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.runner.clone().set_balance(self, address, amount)
self.strategy.runner.set_balance(self, address, amount)
}

/// Gets the balance of an account
Expand All @@ -258,7 +258,7 @@ impl Executor {

/// Set the nonce of an account.
pub fn set_nonce(&mut self, address: Address, nonce: u64) -> BackendResult<()> {
self.strategy.runner.clone().set_nonce(self, address, nonce)
self.strategy.runner.set_nonce(self, address, nonce)
}

/// Returns the nonce of an account.
Expand Down
28 changes: 5 additions & 23 deletions crates/evm/evm/src/executors/strategy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,28 +43,24 @@ impl ExecutorStrategyContext for () {
#[derive(Debug)]
pub struct ExecutorStrategy {
/// Strategy runner.
pub runner: Box<dyn ExecutorStrategyRunner>,
pub runner: &'static dyn ExecutorStrategyRunner,
/// Strategy context.
pub context: Box<dyn ExecutorStrategyContext>,
}

impl ExecutorStrategy {
pub fn new_evm() -> Self {
Self { runner: Box::new(EvmExecutorStrategyRunner::default()), context: Box::new(()) }
Self { runner: &EvmExecutorStrategyRunner, context: Box::new(()) }
}
}

impl Clone for ExecutorStrategy {
fn clone(&self) -> Self {
Self { runner: self.runner.new_cloned(), context: self.context.new_cloned() }
Self { runner: self.runner, context: self.context.new_cloned() }
}
}

pub trait ExecutorStrategyRunner: Debug + Send + Sync + ExecutorStrategyExt {
fn name(&self) -> &'static str;

fn new_cloned(&self) -> Box<dyn ExecutorStrategyRunner>;

fn set_balance(
&self,
executor: &mut Executor,
Expand Down Expand Up @@ -140,17 +136,9 @@ pub trait ExecutorStrategyExt {

/// Implements [ExecutorStrategyRunner] for EVM.
#[derive(Debug, Default, Clone)]
pub struct EvmExecutorStrategyRunner {}
pub struct EvmExecutorStrategyRunner;

impl ExecutorStrategyRunner for EvmExecutorStrategyRunner {
fn name(&self) -> &'static str {
"evm"
}

fn new_cloned(&self) -> Box<dyn ExecutorStrategyRunner> {
Box::new(self.clone())
}

fn set_balance(
&self,
executor: &mut Executor,
Expand Down Expand Up @@ -209,16 +197,10 @@ impl ExecutorStrategyRunner for EvmExecutorStrategyRunner {
_ctx: &dyn ExecutorStrategyContext,
) -> CheatcodeInspectorStrategy {
CheatcodeInspectorStrategy {
runner: Box::new(EvmCheatcodeInspectorStrategyRunner::default()),
runner: &EvmCheatcodeInspectorStrategyRunner,
context: Box::new(()),
}
}
}

impl ExecutorStrategyExt for EvmExecutorStrategyRunner {}

impl Clone for Box<dyn ExecutorStrategyRunner> {
fn clone(&self) -> Self {
self.new_cloned()
}
}
20 changes: 5 additions & 15 deletions crates/strategy/zksync/src/backend.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,19 +66,9 @@ impl BackendStrategyContext for ZksyncBackendStrategyContext {

/// ZKsync implementation for [BackendStrategyRunner].
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct ZksyncBackendStrategyRunner {
evm: EvmBackendStrategyRunner,
}
pub struct ZksyncBackendStrategyRunner;

impl BackendStrategyRunner for ZksyncBackendStrategyRunner {
fn name(&self) -> &'static str {
"zk"
}

fn new_cloned(&self) -> Box<dyn BackendStrategyRunner> {
Box::new(self.clone())
}

fn inspect(
&self,
backend: &mut Backend,
Expand All @@ -87,7 +77,7 @@ impl BackendStrategyRunner for ZksyncBackendStrategyRunner {
inspect_ctx: Box<dyn Any>,
) -> Result<ResultAndState> {
if !is_zksync_inspect_context(inspect_ctx.as_ref()) {
return self.evm.inspect(backend, env, inspector, inspect_ctx);
return EvmBackendStrategyRunner.inspect(backend, env, inspector, inspect_ctx);
}

let inspect_ctx = get_inspect_context(inspect_ctx);
Expand Down Expand Up @@ -137,7 +127,7 @@ impl BackendStrategyRunner for ZksyncBackendStrategyRunner {
active_journaled_state: &JournaledState,
fork_journaled_state: &mut JournaledState,
) {
self.evm.merge_journaled_state_data(
EvmBackendStrategyRunner.merge_journaled_state_data(
ctx,
addr,
active_journaled_state,
Expand All @@ -161,7 +151,7 @@ impl BackendStrategyRunner for ZksyncBackendStrategyRunner {
active: &ForkDB,
fork_db: &mut ForkDB,
) {
self.evm.merge_db_account_data(ctx, addr, active, fork_db);
EvmBackendStrategyRunner.merge_db_account_data(ctx, addr, active, fork_db);
let ctx = get_context(ctx);
let zk_state =
&ZksyncMergeState { persistent_immutable_keys: &ctx.persistent_immutable_keys };
Expand Down Expand Up @@ -235,7 +225,7 @@ pub trait ZksyncBackendStrategyBuilder {
impl ZksyncBackendStrategyBuilder for BackendStrategy {
fn new_zksync() -> Self {
Self {
runner: Box::new(ZksyncBackendStrategyRunner::default()),
runner: &ZksyncBackendStrategyRunner,
context: Box::new(ZksyncBackendStrategyContext::default()),
}
}
Expand Down
18 changes: 4 additions & 14 deletions crates/strategy/zksync/src/cheatcode.rs
Original file line number Diff line number Diff line change
Expand Up @@ -91,9 +91,7 @@ macro_rules! bail {

/// ZKsync implementation for [CheatcodeInspectorStrategyRunner].
#[derive(Debug, Default, Clone)]
pub struct ZksyncCheatcodeInspectorStrategyRunner {
evm: EvmCheatcodeInspectorStrategyRunner,
}
pub struct ZksyncCheatcodeInspectorStrategyRunner;

/// Context for [ZksyncCheatcodeInspectorStrategyRunner].
#[derive(Debug, Default, Clone)]
Expand Down Expand Up @@ -246,14 +244,6 @@ impl ZkPersistNonceUpdate {
}

impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner {
fn name(&self) -> &'static str {
"zk"
}

fn new_cloned(&self) -> Box<dyn CheatcodeInspectorStrategyRunner> {
Box::new(self.clone())
}

fn base_contract_deployed(&self, ctx: &mut dyn CheatcodeInspectorStrategyContext) {
let ctx = get_context(ctx);

Expand Down Expand Up @@ -565,7 +555,7 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
) {
let ctx_zk = get_context(ctx);
if !ctx_zk.using_zk_vm {
return self.evm.record_broadcastable_create_transactions(
return EvmCheatcodeInspectorStrategyRunner.record_broadcastable_create_transactions(
ctx,
config,
input,
Expand Down Expand Up @@ -692,7 +682,7 @@ impl CheatcodeInspectorStrategyRunner for ZksyncCheatcodeInspectorStrategyRunner
let ctx_zk = get_context(ctx);

if !ctx_zk.using_zk_vm {
return self.evm.record_broadcastable_call_transactions(
return EvmCheatcodeInspectorStrategyRunner.record_broadcastable_call_transactions(
ctx,
config,
call,
Expand Down Expand Up @@ -1536,7 +1526,7 @@ pub trait ZksyncCheatcodeInspectorStrategyBuilder {
impl ZksyncCheatcodeInspectorStrategyBuilder for CheatcodeInspectorStrategy {
fn new_zksync(dual_compiled_contracts: DualCompiledContracts, zk_env: ZkEnv) -> Self {
Self {
runner: Box::new(ZksyncCheatcodeInspectorStrategyRunner::default()),
runner: &ZksyncCheatcodeInspectorStrategyRunner,
context: Box::new(ZksyncCheatcodeInspectorStrategyContext::new(
dual_compiled_contracts,
zk_env,
Expand Down
Loading
Loading