From e498b3477b9f23b8a8d11aa591e96efce5c0370f Mon Sep 17 00:00:00 2001 From: Juan Rigada <62958725+Jrigada@users.noreply.github.com> Date: Tue, 22 Oct 2024 12:10:34 -0300 Subject: [PATCH] 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 --- crates/cheatcodes/src/fs.rs | 33 ++++++++++++++------- crates/forge/tests/it/zk/cheats.rs | 11 +++++++ crates/zksync/compiler/src/zksolc/mod.rs | 37 ++++++++++++++++++++++-- testdata/zk/Cheatcodes.t.sol | 19 ++++++++++++ 4 files changed, 87 insertions(+), 13 deletions(-) diff --git a/crates/cheatcodes/src/fs.rs b/crates/cheatcodes/src/fs.rs index da6909401..07c15aa06 100644 --- a/crates/cheatcodes/src/fs.rs +++ b/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 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::>(); - (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")) } diff --git a/crates/forge/tests/it/zk/cheats.rs b/crates/forge/tests/it/zk/cheats.rs index ca33723a0..efc903d96 100644 --- a/crates/forge/tests/it/zk/cheats.rs +++ b/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(); diff --git a/crates/zksync/compiler/src/zksolc/mod.rs b/crates/zksync/compiler/src/zksolc/mod.rs index db4c9b6e1..704b228c0 100644 --- a/crates/zksync/compiler/src/zksolc/mod.rs +++ b/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, + /// 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 { + 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 { self.contracts.iter() diff --git a/testdata/zk/Cheatcodes.t.sol b/testdata/zk/Cheatcodes.t.sol index d19736726..c527a048e 100644 --- a/testdata/zk/Cheatcodes.t.sol +++ b/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 {