feat: Add support for vm.getCode in Zk context

crates/cheatcodes/src/fs.rs

@@ -10,6 +10,7 @@ use dialoguer::{Input, Password};
 use foundry_common::fs;
 use foundry_config::fs_permissions::FsAccessKind;
 use foundry_evm_core::backend::DatabaseExt;
+use foundry_zksync_compiler::ContractType;
 use revm::interpreter::CreateInputs;
 use semver::Version;
 use std::{
@@ -393,17 +394,27 @@ fn get_artifact_code(state: &Cheatcodes, path: &str, deployed: bool) -> Result<B
                 [] => Err(fmt_err!("no matching artifact found")),
                 [artifact] => Ok(artifact),
                 filtered => {
-                    // If we know the current script/test contract solc version, try to filter by it
-                    state
-                        .config
-                        .running_version
-                        .as_ref()
-                        .and_then(|version| {
-                            let filtered = filtered
-                                .iter()
-                                .filter(|(id, _)| id.version == *version)
-                                .collect::<Vec<_>>();
-                            (filtered.len() == 1).then(|| filtered[0])
+                    // If we find more than one artifact, we need to filter by contract type
+                    // depending on whether we are using the zkvm or evm
+                    filtered
+                        .iter()
+                        .find(|(id, _)| {
+                            let contract_type = state
+                                .config
+                                .dual_compiled_contracts
+                                .get_contract_type_by_artifact(id);
+                            match contract_type {
+                                Some(ContractType::ZK) => state.use_zk_vm,
+                                Some(ContractType::EVM) => !state.use_zk_vm,
+                                None => false,
+                            }
+                        })
+                        .or_else(|| {
+                            // If we know the current script/test contract solc version, try to
+                            // filter by it
+                            state.config.running_version.as_ref().and_then(|version| {
+                                filtered.iter().find(|(id, _)| id.version == *version)
+                            })
                         })
                         .ok_or_else(|| fmt_err!("multiple matching artifacts found"))
                 }

crates/forge/tests/it/zk/cheats.rs

@@ -13,6 +13,26 @@ async fn test_zk_cheat_roll_works() {
     TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await;
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn test_zk_cheat_get_code() {
+    let mut zk_config = TEST_DATA_DEFAULT.zk_test_data.zk_config.clone();
+    zk_config.fs_permissions.add(PathPermission::read_write("./zk/zkout/ConstantNumber.sol"));
+    let runner = TEST_DATA_DEFAULT.runner_with_zksync_config(zk_config);
+    let filter = Filter::new("testZkCheatcodesGetCode", "ZkCheatcodesTest", ".*");
+
+    TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_zk_cheat_get_code_evm() {
+    let mut zk_config = TEST_DATA_DEFAULT.zk_test_data.zk_config.clone();
+    zk_config.fs_permissions.add(PathPermission::read_write("./zk/zkout/ConstantNumber.sol"));
+    let runner = TEST_DATA_DEFAULT.runner_with_zksync_config(zk_config);
+    let filter = Filter::new("testZkCheatcodesGetCodeEVM", "ZkCheatcodesTest", ".*");
+
+    TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await;
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn test_zk_cheat_warp_works() {
     let runner = TEST_DATA_DEFAULT.runner_zksync();

crates/zksync/compiler/src/zksolc/mod.rs

@@ -1,18 +1,29 @@
 //! ZKSolc module.
 use std::{
     collections::{HashMap, HashSet, VecDeque},
+    path::PathBuf,
     str::FromStr,
 };
 
 use foundry_compilers::{
     solc::SolcLanguage, zksync::compile::output::ProjectCompileOutput as ZkProjectCompileOutput,
-    Artifact, ArtifactOutput, ConfigurableArtifacts, ProjectCompileOutput, ProjectPathsConfig,
+    Artifact, ArtifactId, ArtifactOutput, ConfigurableArtifacts, ProjectCompileOutput,
+    ProjectPathsConfig,
 };
 
 use alloy_primitives::{keccak256, B256};
 use tracing::debug;
 use zksync_types::H256;
 
+/// Represents the type of contract (ZK or EVM)
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum ContractType {
+    /// ZkSolc compiled contract
+    ZK,
+    /// Solc compiled contract
+    EVM,
+}
+
 /// Defines a contract that has been dual compiled with both zksolc and solc
 #[derive(Debug, Default, Clone)]
 pub struct DualCompiledContract {
@@ -36,6 +47,10 @@ pub struct DualCompiledContract {
 #[derive(Debug, Default, Clone)]
 pub struct DualCompiledContracts {
     contracts: Vec<DualCompiledContract>,
+    /// ZKvm artifacts path
+    pub zk_artifact_path: PathBuf,
+    /// EVM artifacts path
+    pub evm_artifact_path: PathBuf,
 }
 
 impl DualCompiledContracts {
@@ -158,7 +173,11 @@ impl DualCompiledContracts {
             }
         }
 
-        Self { contracts: dual_compiled_contracts }
+        Self {
+            contracts: dual_compiled_contracts,
+            zk_artifact_path: zk_layout.artifacts.clone(),
+            evm_artifact_path: layout.artifacts.clone(),
+        }
     }
 
     /// Finds a contract matching the ZK deployed bytecode
@@ -209,6 +228,23 @@ impl DualCompiledContracts {
         visited.into_iter().cloned().collect()
     }
 
+    /// Returns the contract type (ZK or EVM) based on the artifact path
+    pub fn get_contract_type_by_artifact(&self, artifact_id: &ArtifactId) -> Option<ContractType> {
+        if artifact_id.path.starts_with(&self.zk_artifact_path) {
+            Some(ContractType::ZK)
+        } else if artifact_id.path.starts_with(&self.evm_artifact_path) {
+            Some(ContractType::EVM)
+        } else {
+            tracing::error!(
+                "Contract path {:?} not found in {:?} or {:?}",
+                artifact_id.path,
+                self.zk_artifact_path,
+                self.evm_artifact_path
+            );
+            None
+        }
+    }
+
     /// Returns an iterator over all `[DualCompiledContract]`s in the collection
     pub fn iter(&self) -> impl Iterator<Item = &DualCompiledContract> {
         self.contracts.iter()

testdata/zk/Cheatcodes.t.sol

@@ -127,6 +127,30 @@ contract ZkCheatcodesTest is DSTest {
         require(vm.getNonce(TEST_ADDRESS) == 0, "era nonce mismatch");
     }
 
+    function testZkCheatcodesGetCode() public {
+        bytes memory fullPath = vm.getCode("ConstantNumber.sol");
+
+        string memory expected = string(
+            bytes(
+                hex"0000008003000039000000400030043f0000000100200190000000120000c13d000000000201001900000009002001980000001a0000613d000000000101043b0000000a011001970000000b0010009c0000001a0000c13d0000000001000416000000000001004b0000001a0000c13d0000000a01000039000000800010043f0000000c010000410000001d0001042e0000000001000416000000000001004b0000001a0000c13d00000020010000390000010000100443000001200000044300000008010000410000001d0001042e00000000010000190000001e000104300000001c000004320000001d0001042e0000001e000104300000000000000000000000020000000000000000000000000000004000000100000000000000000000000000000000000000000000000000fffffffc000000000000000000000000ffffffff00000000000000000000000000000000000000000000000000000000643ceff90000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000200000008000000000000000000000000000000000000000000000000000000000000000000000000000000000579fa5ae1b1bfd1afd3813b5228dda992ad50456891069ef9c5d8a204d0f0487"
+            )
+        );
+        assertEq(string(fullPath), expected, "code for the contract was incorrect");
+    }
+
+    function testZkCheatcodesGetCodeEVM() public {
+        vm.zkVm(false);
+        bytes memory fullPath = vm.getCode("ConstantNumber.sol");
+        vm.zkVm(true);
+
+        string memory expected = string(
+            bytes(
+                hex"6080604052348015600f57600080fd5b50607780601d6000396000f3fe6080604052348015600f57600080fd5b506004361060285760003560e01c8063643ceff914602d575b600080fd5b60408051600a815290519081900360200190f3fea2646970667358221220e940a7e80d99723067bd50f2ae220c5faa1bc6a1d90f422384a14c93135a110064736f6c634300081b0033"
+            )
+        );
+        assertEq(string(fullPath), expected, "code for the contract was incorrect");
+    }
+
     function testZkCheatcodesEtch() public {
         vm.selectFork(forkEra);