remove use_zk completely, fix ExecutorStrategy trait

crates/cheatcodes/src/strategy.rs

@@ -19,6 +19,8 @@ use crate::{
 };
 
 pub trait CheatcodeInspectorStrategy: Debug + Send + Sync {
+    fn name(&self) -> &'static str;
+
     /// Get nonce.
     fn get_nonce(&mut self, ccx: &mut CheatsCtxt, address: Address) -> Result<u64> {
         let account = ccx.ecx.journaled_state.load_account(address, &mut ccx.ecx.db)?;
@@ -179,23 +181,23 @@ pub trait CheatcodeInspectorStrategy: Debug + Send + Sync {
 
 /// We define this in our fork
 pub trait CheatcodeInspectorStrategyExt: CheatcodeInspectorStrategy {
-    fn zksync_skip_zkvm(&mut self) -> Result {
+    fn zksync_cheatcode_skip_zkvm(&mut self) -> Result {
         unimplemented!()
     }
 
-    fn zksync_set_paymaster(
+    fn zksync_cheatcode_set_paymaster(
         &mut self,
         _paymaster_address: Address,
         _paymaster_input: &Bytes,
     ) -> Result {
         unimplemented!()
     }
 
-    fn zksync_use_factory_deps(&mut self, _name: String) -> Result {
+    fn zksync_cheatcode_use_factory_deps(&mut self, _name: String) -> Result {
         unimplemented!()
     }
 
-    fn zksync_register_contract(
+    fn zksync_cheatcode_register_contract(
         &mut self,
         _name: String,
         _zk_bytecode_hash: FixedBytes<32>,
@@ -208,17 +210,15 @@ pub trait CheatcodeInspectorStrategyExt: CheatcodeInspectorStrategy {
         unimplemented!()
     }
 
-    fn zksync_record_create_address(&mut self, _outcome: &CreateOutcome) {
+    fn zksync_cheatcode_select_zk_vm(&mut self, _data: InnerEcx, _enable: bool) {
         unimplemented!()
     }
 
-    fn zksync_sync_nonce(&mut self, _sender: Address, _nonce: u64, _ecx: Ecx) {
-        unimplemented!()
-    }
+    fn zksync_record_create_address(&mut self, _outcome: &CreateOutcome) {}
 
-    fn zksync_set_deployer_call_input(&mut self, _call: &mut CallInputs) {
-        unimplemented!()
-    }
+    fn zksync_sync_nonce(&mut self, _sender: Address, _nonce: u64, _ecx: Ecx) {}
+
+    fn zksync_set_deployer_call_input(&mut self, _call: &mut CallInputs) {}
 
     fn zksync_try_create(
         &mut self,
@@ -227,7 +227,7 @@ pub trait CheatcodeInspectorStrategyExt: CheatcodeInspectorStrategy {
         _input: &dyn CommonCreateInput,
         _executor: &mut dyn CheatcodesExecutor,
     ) -> Option<CreateOutcome> {
-        unimplemented!()
+        None
     }
 
     fn zksync_try_call(
@@ -237,22 +237,20 @@ pub trait CheatcodeInspectorStrategyExt: CheatcodeInspectorStrategy {
         _input: &CallInputs,
         _executor: &mut dyn CheatcodesExecutor,
     ) -> Option<CallOutcome> {
-        unimplemented!()
+        None
     }
 
-    fn zksync_select_fork_vm(&mut self, _data: InnerEcx, _fork_id: LocalForkId) {
-        unimplemented!()
-    }
-
-    fn zksync_select_zk_vm(&mut self, _data: InnerEcx, _enable: bool) {
-        unimplemented!()
-    }
+    fn zksync_select_fork_vm(&mut self, _data: InnerEcx, _fork_id: LocalForkId) {}
 }
 
 #[derive(Debug, Default, Clone)]
 pub struct EvmCheatcodeInspectorStrategy;
 
 impl CheatcodeInspectorStrategy for EvmCheatcodeInspectorStrategy {
+    fn name(&self) -> &'static str {
+        "evm"
+    }
+
     fn record_broadcastable_create_transactions(
         &mut self,
         _config: Arc<CheatsConfig>,

crates/cheatcodes/src/test.rs

@@ -19,15 +19,15 @@ impl Cheatcode for zkVmCall {
             .strategy
             .lock()
             .expect("failed acquiring strategy")
-            .zksync_select_zk_vm(ccx.ecx, enable);
+            .zksync_cheatcode_select_zk_vm(ccx.ecx, enable);
 
         Ok(Default::default())
     }
 }
 
 impl Cheatcode for zkVmSkipCall {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
-        ccx.state.strategy.lock().expect("failed acquiring strategy").zksync_skip_zkvm()
+        ccx.state.strategy.lock().expect("failed acquiring strategy").zksync_cheatcode_skip_zkvm()
     }
 }
 
@@ -38,7 +38,7 @@ impl Cheatcode for zkUsePaymasterCall {
             .strategy
             .lock()
             .expect("failed acquiring strategy")
-            .zksync_set_paymaster(*paymaster_address, paymaster_input)
+            .zksync_cheatcode_set_paymaster(*paymaster_address, paymaster_input)
     }
 }
 
@@ -49,7 +49,7 @@ impl Cheatcode for zkUseFactoryDepCall {
             .strategy
             .lock()
             .expect("failed acquiring strategy")
-            .zksync_use_factory_deps(name.clone())
+            .zksync_cheatcode_use_factory_deps(name.clone())
     }
 }
 
@@ -64,15 +64,19 @@ impl Cheatcode for zkRegisterContractCall {
             zkDeployedBytecode,
         } = self;
 
-        ccx.state.strategy.lock().expect("failed acquiring strategy").zksync_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
+            .lock()
+            .expect("failed acquiring 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(),
+            )
     }
 }
 

crates/chisel/src/executor.rs

@@ -15,7 +15,7 @@ use foundry_evm::{
     backend::Backend,
     decode::decode_console_logs,
     executors::{
-        strategy::{EvmExecutorStrategy, ExecutorStrategy},
+        strategy::{EvmExecutorStrategy, ExecutorStrategyExt},
         ExecutorBuilder,
     },
     inspectors::CheatsConfig,
@@ -325,7 +325,7 @@ impl SessionSource {
         let env =
             self.config.evm_opts.evm_env().await.expect("Could not instantiate fork environment");
 
-        let executor_strategy: Arc<Mutex<dyn ExecutorStrategy>> =
+        let executor_strategy: Arc<Mutex<dyn ExecutorStrategyExt>> =
             if self.config.foundry_config.zksync.run_in_zk_mode() {
                 Arc::new(Mutex::new(ZksyncExecutorStrategy::default()))
             } else {
@@ -362,7 +362,7 @@ impl SessionSource {
                         executor_strategy
                             .lock()
                             .expect("failed acquiring strategy")
-                            .new_cheatcode_inspector_strategy(Default::default()),
+                            .new_cheatcode_inspector_strategy(),
                     )
                     .into(),
                 )

crates/cli/src/utils/mod.rs

@@ -9,7 +9,7 @@ use foundry_common::{
     shell,
 };
 use foundry_config::{Chain, Config};
-use foundry_evm::executors::strategy::{EvmExecutorStrategy, ExecutorStrategy};
+use foundry_evm::executors::strategy::{EvmExecutorStrategy, ExecutorStrategyExt};
 use foundry_strategy_zksync::ZksyncExecutorStrategy;
 use serde::de::DeserializeOwned;
 use std::{
@@ -94,7 +94,7 @@ pub fn get_provider(config: &Config) -> Result<RetryProvider> {
     get_provider_builder(config)?.build()
 }
 
-pub fn get_executor_strategy(config: &Config) -> Arc<Mutex<dyn ExecutorStrategy>> {
+pub fn get_executor_strategy(config: &Config) -> Arc<Mutex<dyn ExecutorStrategyExt>> {
     if config.zksync.run_in_zk_mode() {
         Arc::new(Mutex::new(ZksyncExecutorStrategy::default()))
     } else {

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::{BackendError, ForkInfo};
+use super::{strategy::BackendStrategyExt, BackendError, ForkInfo};
 use crate::{
     backend::{
         diagnostic::RevertDiagnostic, Backend, DatabaseExt, LocalForkId, RevertStateSnapshotAction,
@@ -22,7 +22,8 @@ use revm::{
 };
 use std::{
     borrow::Cow,
-    collections::{BTreeMap, HashSet},
+    collections::BTreeMap,
+    sync::{Arc, Mutex},
 };
 
 /// A wrapper around `Backend` that ensures only `revm::DatabaseRef` functions are called.
@@ -118,8 +119,8 @@ impl DatabaseExt for CowBackend<'_> {
         self.backend.to_mut().get_fork_info(id)
     }
 
-    fn save_zk_immutable_storage(&mut self, addr: Address, keys: HashSet<U256>) {
-        self.backend.to_mut().save_zk_immutable_storage(addr, keys)
+    fn get_strategy(&mut self) -> Arc<Mutex<dyn BackendStrategyExt>> {
+        self.backend.as_ref().strategy.clone()
     }
 
     fn snapshot_state(&mut self, journaled_state: &JournaledState, env: &Env) -> U256 {

crates/evm/core/src/backend/mod.rs

@@ -31,7 +31,7 @@ use std::{
     sync::{Arc, Mutex},
     time::Instant,
 };
-use strategy::{BackendStrategy, BackendStrategyForkInfo};
+use strategy::{BackendStrategyExt, BackendStrategyForkInfo};
 
 mod diagnostic;
 pub use diagnostic::RevertDiagnostic;
@@ -102,13 +102,8 @@ pub trait DatabaseExt: Database<Error = DatabaseError> + DatabaseCommit {
     /// and the the fork environment.
     fn get_fork_info(&mut self, id: LocalForkId) -> eyre::Result<ForkInfo>;
 
-    /// Saves the storage keys for immutable variables per address.
-    ///
-    /// These are required during fork to help merge the persisted addresses, as they are stored
-    /// 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 save_zk_immutable_storage(&mut self, addr: Address, keys: HashSet<U256>);
+    /// Retrieve the strategy.
+    fn get_strategy(&mut self) -> Arc<Mutex<dyn BackendStrategyExt>>;
 
     /// Reverts the snapshot if it exists
     ///
@@ -466,7 +461,7 @@ struct _ObjectSafe(dyn DatabaseExt);
 #[must_use]
 pub struct Backend {
     /// The behavior strategy.
-    pub strategy: Arc<Mutex<dyn BackendStrategy>>,
+    pub strategy: Arc<Mutex<dyn BackendStrategyExt>>,
 
     /// The access point for managing forks
     forks: MultiFork,
@@ -497,23 +492,14 @@ pub struct Backend {
     inner: BackendInner,
     /// Keeps track of the fork type
     fork_url_type: CachedForkType,
-    /// TODO: Ensure this parameter is updated on `select_fork`.
-    ///
-    /// Keeps track if the backend is in ZK mode.
-    /// This is required to correctly merge storage when selecting another ZK fork.
-    /// The balance, nonce and code are stored under zkSync's respective system contract
-    /// storages. These need to be merged into the forked storage.
-    pub is_zk: bool,
-    /// Store storage keys per contract address for immutable variables.
-    zk_recorded_immutable_keys: HashMap<Address, HashSet<U256>>,
 }
 
 impl Backend {
     /// Creates a new Backend with a spawned multi fork thread.
     ///
     /// If `fork` is `Some` this will use a `fork` database, otherwise with an in-memory
     /// database.
-    pub fn spawn(fork: Option<CreateFork>, strategy: Arc<Mutex<dyn BackendStrategy>>) -> Self {
+    pub fn spawn(fork: Option<CreateFork>, strategy: Arc<Mutex<dyn BackendStrategyExt>>) -> Self {
         Self::new(MultiFork::spawn(), fork, strategy)
     }
 
@@ -526,7 +512,7 @@ impl Backend {
     pub fn new(
         forks: MultiFork,
         fork: Option<CreateFork>,
-        strategy: Arc<Mutex<dyn BackendStrategy>>,
+        strategy: Arc<Mutex<dyn BackendStrategyExt>>,
     ) -> Self {
         trace!(target: "backend", forking_mode=?fork.is_some(), "creating executor backend");
         // Note: this will take of registering the `fork`
@@ -542,8 +528,6 @@ impl Backend {
             active_fork_ids: None,
             inner,
             fork_url_type: Default::default(),
-            is_zk: false,
-            zk_recorded_immutable_keys: Default::default(),
             strategy,
         };
 
@@ -571,7 +555,7 @@ impl Backend {
         id: &ForkId,
         fork: Fork,
         journaled_state: JournaledState,
-        strategy: Arc<Mutex<dyn BackendStrategy>>,
+        strategy: Arc<Mutex<dyn BackendStrategyExt>>,
     ) -> Self {
         let mut backend = Self::spawn(None, strategy);
         let fork_ids = backend.inner.insert_new_fork(id.clone(), fork.db, journaled_state);
@@ -589,8 +573,6 @@ impl Backend {
             active_fork_ids: None,
             inner: Default::default(),
             fork_url_type: Default::default(),
-            is_zk: false,
-            zk_recorded_immutable_keys: Default::default(),
             strategy: self.strategy.clone(),
         }
     }
@@ -996,11 +978,8 @@ impl DatabaseExt for Backend {
         Ok(ForkInfo { fork_type, fork_env })
     }
 
-    fn save_zk_immutable_storage(&mut self, addr: Address, keys: HashSet<U256>) {
-        self.zk_recorded_immutable_keys
-            .entry(addr)
-            .and_modify(|entry| entry.extend(&keys))
-            .or_insert(keys);
+    fn get_strategy(&mut self) -> Arc<Mutex<dyn BackendStrategyExt>> {
+        self.strategy.clone()
     }
 
     fn snapshot_state(&mut self, journaled_state: &JournaledState, env: &Env) -> U256 {
@@ -1875,7 +1854,7 @@ impl BackendInner {
         id: LocalForkId,
         new_fork_id: ForkId,
         backend: SharedBackend,
-        strategy: Arc<Mutex<dyn BackendStrategy>>,
+        strategy: Arc<Mutex<dyn BackendStrategyExt>>,
     ) -> eyre::Result<ForkLookupIndex> {
         let fork_id = self.ensure_fork_id(id)?;
         let idx = self.ensure_fork_index(fork_id)?;
@@ -2007,7 +1986,7 @@ fn commit_transaction(
     fork_id: &ForkId,
     persistent_accounts: &HashSet<Address>,
     inspector: &mut dyn InspectorExt,
-    strategy: Arc<Mutex<dyn BackendStrategy>>,
+    strategy: Arc<Mutex<dyn BackendStrategyExt>>,
 ) -> 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