use Box instead of Arc<Mutex<T>>

crates/cheatcodes/src/config.rs

@@ -14,7 +14,6 @@ use foundry_evm_core::opts::EvmOpts;
 use semver::Version;
 use std::{
     path::{Path, PathBuf},
-    sync::{Arc, Mutex},
     time::Duration,
 };
 
@@ -58,7 +57,7 @@ pub struct CheatsConfig {
     /// Version of the script/test contract which is currently running.
     pub running_version: Option<Version>,
     /// The behavior strategy.
-    pub strategy: Arc<Mutex<dyn CheatcodeInspectorStrategyExt>>,
+    pub strategy: Box<dyn CheatcodeInspectorStrategyExt>,
     /// Whether to enable legacy (non-reverting) assertions.
     pub assertions_revert: bool,
     /// Optional seed for the RNG algorithm.
@@ -74,7 +73,7 @@ impl CheatsConfig {
         available_artifacts: Option<ContractsByArtifact>,
         running_contract: Option<String>,
         running_version: Option<Version>,
-        strategy: Arc<Mutex<dyn CheatcodeInspectorStrategyExt>>,
+        strategy: Box<dyn CheatcodeInspectorStrategyExt>,
     ) -> Self {
         let mut allowed_paths = vec![config.root.0.clone()];
         allowed_paths.extend(config.libs.clone());
@@ -235,7 +234,7 @@ impl Default for CheatsConfig {
             available_artifacts: Default::default(),
             running_contract: Default::default(),
             running_version: Default::default(),
-            strategy: Arc::new(Mutex::new(EvmCheatcodeInspectorStrategy::default())),
+            strategy: Box::new(EvmCheatcodeInspectorStrategy::default()),
             assertions_revert: true,
             seed: None,
         }
@@ -254,7 +253,7 @@ mod tests {
             None,
             None,
             None,
-            Arc::new(Mutex::new(EvmCheatcodeInspectorStrategy::default())),
+            Box::new(EvmCheatcodeInspectorStrategy::default()),
         )
     }
 

crates/cheatcodes/src/evm.rs

@@ -64,9 +64,7 @@ impl Cheatcode for getNonce_0Call {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { account } = self;
 
-        let strategy = ccx.state.strategy.clone();
-        let mut guard = strategy.lock().expect("failed acquiring strategy");
-        guard.cheatcode_get_nonce(ccx, *account)
+        ccx.with_strategy(|strategy, ccx| strategy.cheatcode_get_nonce(ccx, *account))
     }
 }
 
@@ -352,9 +350,7 @@ impl Cheatcode for getBlobhashesCall {
 impl Cheatcode for rollCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { newHeight } = self;
-        let strategy = ccx.state.strategy.clone();
-        let mut guard = strategy.lock().expect("failed acquiring strategy");
-        guard.cheatcode_roll(ccx, *newHeight)
+        ccx.with_strategy(|strategy, ccx| strategy.cheatcode_roll(ccx, *newHeight))
     }
 }
 
@@ -376,9 +372,7 @@ impl Cheatcode for txGasPriceCall {
 impl Cheatcode for warpCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { newTimestamp } = self;
-        let strategy = ccx.state.strategy.clone();
-        let mut guard = strategy.lock().expect("failed acquiring strategy");
-        guard.cheatcode_warp(ccx, *newTimestamp)
+        ccx.with_strategy(|strategy, ccx| strategy.cheatcode_warp(ccx, *newTimestamp))
     }
 }
 
@@ -413,48 +407,40 @@ impl Cheatcode for dealCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { account: address, newBalance: new_balance } = *self;
 
-        let strategy = ccx.state.strategy.clone();
-        let mut guard = strategy.lock().expect("failed acquiring strategy");
-        guard.cheatcode_deal(ccx, address, new_balance)
+        ccx.with_strategy(|strategy, ccx| strategy.cheatcode_deal(ccx, address, new_balance))
     }
 }
 
 impl Cheatcode for etchCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { target, newRuntimeBytecode } = self;
 
-        let strategy = ccx.state.strategy.clone();
-        let mut guard = strategy.lock().expect("failed acquiring strategy");
-        guard.cheatcode_etch(ccx, *target, newRuntimeBytecode)
+        ccx.with_strategy(|strategy, ccx| strategy.cheatcode_etch(ccx, *target, newRuntimeBytecode))
     }
 }
 
 impl Cheatcode for resetNonceCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { account } = self;
-        let strategy = ccx.state.strategy.clone();
-        let mut guard = strategy.lock().expect("failed acquiring strategy");
-        guard.cheatcode_reset_nonce(ccx, *account)
+        ccx.with_strategy(|strategy, ccx| strategy.cheatcode_reset_nonce(ccx, *account))
     }
 }
 
 impl Cheatcode for setNonceCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { account, newNonce } = *self;
 
-        let strategy = ccx.state.strategy.clone();
-        let mut guard = strategy.lock().expect("failed acquiring strategy");
-        guard.cheatcode_set_nonce(ccx, account, newNonce)
+        ccx.with_strategy(|strategy, ccx| strategy.cheatcode_set_nonce(ccx, account, newNonce))
     }
 }
 
 impl Cheatcode for setNonceUnsafeCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { account, newNonce } = *self;
 
-        let strategy = ccx.state.strategy.clone();
-        let mut guard = strategy.lock().expect("failed acquiring strategy");
-        guard.cheatcode_set_nonce_unsafe(ccx, account, newNonce)
+        ccx.with_strategy(|strategy, ccx| {
+            strategy.cheatcode_set_nonce_unsafe(ccx, account, newNonce)
+        })
     }
 }
 

crates/cheatcodes/src/evm/fork.rs

@@ -125,11 +125,7 @@ impl Cheatcode for selectForkCall {
         persist_caller(ccx);
         check_broadcast(ccx.state)?;
 
-        ccx.state
-            .strategy
-            .lock()
-            .expect("failed acquiring strategy")
-            .zksync_select_fork_vm(ccx.ecx, *forkId);
+        ccx.with_strategy(|strategy, ccx| strategy.zksync_select_fork_vm(ccx.ecx, *forkId));
         ccx.ecx.db.select_fork(*forkId, &mut ccx.ecx.env, &mut ccx.ecx.journaled_state)?;
 
         Ok(Default::default())
@@ -285,11 +281,7 @@ 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
-        .lock()
-        .expect("failed acquiring strategy")
-        .zksync_select_fork_vm(ccx.ecx, id);
+    ccx.with_strategy(|strategy, ccx| strategy.zksync_select_fork_vm(ccx.ecx, id));
     ccx.ecx.db.select_fork(id, &mut ccx.ecx.env, &mut ccx.ecx.journaled_state)?;
     Ok(id.abi_encode())
 }

crates/cheatcodes/src/evm/mock.rs

@@ -15,9 +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;
-        let strategy = ccx.state.strategy.clone();
-        let mut guard = strategy.lock().expect("failed acquiring strategy");
-        guard.cheatcode_mock_call(ccx, *callee, data, returnData)
+        ccx.with_strategy(|strategy, ccx| {
+            strategy.cheatcode_mock_call(ccx, *callee, data, returnData)
+        })
     }
 }
 
@@ -85,9 +85,9 @@ impl Cheatcode for mockCalls_1Call {
 impl Cheatcode for mockCallRevert_0Call {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
         let Self { callee, data, revertData } = self;
-        let strategy = ccx.state.strategy.clone();
-        let mut guard = strategy.lock().expect("failed acquiring strategy");
-        guard.cheatcode_mock_call_revert(ccx, *callee, data, revertData)
+        ccx.with_strategy(|strategy, ccx| {
+            strategy.cheatcode_mock_call_revert(ccx, *callee, data, revertData)
+        })
     }
 }
 

crates/cheatcodes/src/fs.rs

@@ -283,9 +283,7 @@ impl Cheatcode for getArtifactPathByDeployedCodeCall {
 impl Cheatcode for getCodeCall {
     fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self { artifactPath: path } = self;
-        let strategy = state.strategy.clone();
-        let guard = strategy.lock().expect("failed acquiring strategy");
-        guard.get_artifact_code(state, path, false)
+        state.with_strategy(|strategy, state| strategy.get_artifact_code(state, path, false))
     }
 }
 

crates/cheatcodes/src/inspector.rs

@@ -61,14 +61,15 @@ use std::{
     io::BufReader,
     ops::Range,
     path::PathBuf,
-    sync::{Arc, Mutex},
+    sync::Arc,
 };
 
 mod utils;
 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 CheatcodeInspectorStrategyExt;
 
 /// Helper trait for obtaining complete [revm::Inspector] instance from mutable reference to
 /// [Cheatcodes].
@@ -527,7 +528,7 @@ pub struct Cheatcodes {
     pub wallets: Option<Wallets>,
 
     /// The behavior strategy.
-    pub strategy: Arc<Mutex<dyn CheatcodeInspectorStrategyExt>>,
+    pub strategy: Option<Box<dyn CheatcodeInspectorStrategyExt>>,
 }
 
 impl Clone for Cheatcodes {
@@ -567,7 +568,7 @@ impl Clone for Cheatcodes {
             arbitrary_storage: self.arbitrary_storage.clone(),
             deprecated: self.deprecated.clone(),
             wallets: self.wallets.clone(),
-            strategy: self.strategy.lock().expect("failed acquiring strategy").new_cloned_ext(),
+            strategy: self.strategy.as_ref().map(|s| s.new_cloned_ext()),
         }
     }
 }
@@ -587,7 +588,7 @@ impl Cheatcodes {
         Self {
             fs_commit: true,
             labels: config.labels.clone(),
-            strategy: config.strategy.clone(),
+            strategy: Some(config.strategy.clone()),
             config,
             block: Default::default(),
             active_delegation: Default::default(),
@@ -762,16 +763,13 @@ impl Cheatcodes {
                 if ecx_inner.journaled_state.depth() == broadcast.depth {
                     input.set_caller(broadcast.new_origin);
 
-                    self.strategy
-                        .lock()
-                        .expect("failed acquiring strategy")
-                        .record_broadcastable_create_transactions(
-                            self.config.clone(),
-                            &input,
-                            ecx_inner,
-                            broadcast,
-                            &mut self.broadcastable_transactions,
-                        );
+                    self.strategy.as_mut().unwrap().record_broadcastable_create_transactions(
+                        self.config.clone(),
+                        &input,
+                        ecx_inner,
+                        broadcast,
+                        &mut self.broadcastable_transactions,
+                    );
 
                     input.log_debug(self, &input.scheme().unwrap_or(CreateScheme::Create));
                 }
@@ -803,9 +801,9 @@ impl Cheatcodes {
             }]);
         }
 
-        let strategy = self.strategy.clone();
-        let mut guard = strategy.lock().expect("failed acquiring strategy");
-        if let Some(result) = guard.zksync_try_create(self, ecx, &input, executor) {
+        if let Some(result) = self.with_strategy(|strategy, cheatcodes| {
+            strategy.zksync_try_create(cheatcodes, ecx, &input, executor)
+        }) {
             return Some(result);
         }
 
@@ -920,7 +918,7 @@ where {
         }
 
         self.strategy
-            .lock()
+            .as_mut()
             .expect("failed acquiring strategy")
             .zksync_record_create_address(&outcome);
 
@@ -966,7 +964,7 @@ where {
             let nonce = prev.saturating_sub(1);
             account.info.nonce = nonce;
             self.strategy
-                .lock()
+                .as_mut()
                 .expect("failed acquiring strategy")
                 .zksync_sync_nonce(sender, nonce, ecx);
 
@@ -1003,7 +1001,7 @@ where {
         }
 
         self.strategy
-            .lock()
+            .as_mut()
             .expect("failed acquiring strategy")
             .zksync_set_deployer_call_input(call);
 
@@ -1140,7 +1138,7 @@ where {
                     }
 
                     self.strategy
-                        .lock()
+                        .as_mut()
                         .expect("failed acquiring strategy")
                         .record_broadcastable_call_transactions(
                             self.config.clone(),
@@ -1222,9 +1220,9 @@ where {
             }]);
         }
 
-        let strategy = self.strategy.clone();
-        let mut guard = strategy.lock().expect("failed acquiring strategy");
-        if let Some(result) = guard.zksync_try_call(self, ecx, call, executor) {
+        if let Some(result) = self.with_strategy(|strategy, cheatcodes| {
+            strategy.zksync_try_call(cheatcodes, ecx, call, executor)
+        }) {
             return Some(result);
         }
 
@@ -1266,6 +1264,17 @@ where {
             None => false,
         }
     }
+
+    pub fn with_strategy<F, R>(&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 {
@@ -1286,7 +1295,7 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes {
         }
 
         self.strategy
-            .lock()
+            .as_mut()
             .expect("failed acquiring strategy")
             .post_initialize_interp(interpreter, ecx);
     }
@@ -1333,11 +1342,7 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes {
 
     #[inline]
     fn step_end(&mut self, interpreter: &mut Interpreter, ecx: Ecx) {
-        if self
-            .strategy
-            .try_lock()
-            .expect("failed acquiring strategy")
-            .pre_step_end(interpreter, ecx)
+        if self.strategy.as_mut().expect("failed acquiring strategy").pre_step_end(interpreter, ecx)
         {
             return;
         }

crates/cheatcodes/src/lib.rs

@@ -17,6 +17,7 @@ extern crate tracing;
 
 use alloy_primitives::Address;
 use foundry_evm_core::backend::DatabaseExt;
+use inspector::Strategy;
 use revm::{ContextPrecompiles, InnerEvmContext};
 use spec::Status;
 
@@ -174,4 +175,15 @@ impl CheatsCtxt<'_, '_, '_, '_> {
     pub(crate) fn is_precompile(&self, address: &Address) -> bool {
         self.precompiles.contains(address)
     }
+
+    pub(crate) fn with_strategy<F, R>(&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
+    }
 }