feat: Add support for vm.getCode in Zk context (#604)

* Add integration for vm.getCode in zkcontext

* Format cheatcode contract

* fix test and clippy

* fix comparing with wrong bytecode

* improve tests

* forge fmt

* change test name

* refactor(test:zk): simplify getCode

* cargo fmt

* panic if can't find contract by type

---------

Co-authored-by: Francesco Dainese <franci.dainese@gmail.com>

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,17 @@ 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("./zk"));
+
+    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_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,20 @@ 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 {
+            panic!(
+                "Unexpected artifact path: {:?}. Not found in ZK path {:?} or EVM path {:?}",
+                artifact_id.path, self.zk_artifact_path, self.evm_artifact_path
+            );
+        }
+    }
+
     /// 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,14 @@ contract ZkCheatcodesTest is DSTest {
         require(vm.getNonce(TEST_ADDRESS) == 0, "era nonce mismatch");
     }
 
+    function testZkCheatcodesGetCode() public {
+        string memory contractName = "ConstantNumber";
+        getCodeCheck(contractName, "zkout");
+
+        vm.zkVm(false);
+        getCodeCheck(contractName, "out");
+    }
+
     function testZkCheatcodesEtch() public {
         vm.selectFork(forkEra);
 
@@ -260,6 +268,17 @@ contract ZkCheatcodesTest is DSTest {
 
         assertEq(zkvmEntries.length, evmEntries.length);
     }
+
+    // Utility function
+    function getCodeCheck(string memory contractName, string memory outDir) internal {
+        bytes memory bytecode = vm.getCode(contractName);
+
+        string memory artifactPath = string.concat("zk/", outDir, "/", contractName, ".sol/", contractName, ".json");
+        string memory artifact = vm.readFile(artifactPath);
+        bytes memory expectedBytecode = vm.parseJsonBytes(artifact, ".bytecode.object");
+
+        assertEq(bytecode, expectedBytecode, "code for the contract was incorrect");
+    }
 }
 
 contract UsesCheatcodes {