getCode + getDeployedCode updates

crates/cheatcodes/Cargo.toml

@@ -42,3 +42,4 @@ k256.workspace = true
 walkdir = "2"
 p256 = "0.13.2"
 thiserror = "1"
+semver = "1"

crates/cheatcodes/src/config.rs

@@ -2,7 +2,7 @@ use super::Result;
 use crate::{script::ScriptWallets, Vm::Rpc};
 use alloy_primitives::Address;
 use foundry_common::fs::normalize_path;
-use foundry_compilers::{utils::canonicalize, ProjectPathsConfig};
+use foundry_compilers::{utils::canonicalize, ArtifactId, ProjectPathsConfig};
 use foundry_config::{
     cache::StorageCachingConfig, fs_permissions::FsAccessKind, Config, FsPermissions,
     ResolvedRpcEndpoints,
@@ -40,11 +40,20 @@ pub struct CheatsConfig {
     pub labels: HashMap<Address, String>,
     /// Script wallets
     pub script_wallets: Option<ScriptWallets>,
+    /// Artifacts which are guaranteed to be fresh (either recompiled or cached).
+    /// If Some, `vm.getDeployedCode` invocations are validated to be in scope of this list.
+    /// If None, no validation is performed.
+    pub available_artifacts: Option<Vec<ArtifactId>>,
 }
 
 impl CheatsConfig {
     /// Extracts the necessary settings from the Config
-    pub fn new(config: &Config, evm_opts: EvmOpts, script_wallets: Option<ScriptWallets>) -> Self {
+    pub fn new(
+        config: &Config,
+        evm_opts: EvmOpts,
+        available_artifacts: Option<Vec<ArtifactId>>,
+        script_wallets: Option<ScriptWallets>,
+    ) -> Self {
         let mut allowed_paths = vec![config.__root.0.clone()];
         allowed_paths.extend(config.libs.clone());
         allowed_paths.extend(config.allow_paths.clone());
@@ -64,6 +73,7 @@ impl CheatsConfig {
             evm_opts,
             labels: config.labels.clone(),
             script_wallets,
+            available_artifacts,
         }
     }
 
@@ -180,6 +190,7 @@ impl Default for CheatsConfig {
             evm_opts: Default::default(),
             labels: Default::default(),
             script_wallets: None,
+            available_artifacts: Default::default(),
         }
     }
 }
@@ -194,6 +205,7 @@ mod tests {
             &Config { __root: PathBuf::from(root).into(), fs_permissions, ..Default::default() },
             Default::default(),
             None,
+            None,
         )
     }
 

crates/cheatcodes/src/fs.rs

@@ -4,12 +4,13 @@ use crate::{Cheatcode, Cheatcodes, Result, Vm::*};
 use alloy_json_abi::ContractObject;
 use alloy_primitives::U256;
 use alloy_sol_types::SolValue;
-use foundry_common::{fs, get_artifact_path};
+use foundry_common::fs;
 use foundry_config::fs_permissions::FsAccessKind;
+use semver::Version;
 use std::{
     collections::hash_map::Entry,
     io::{BufRead, BufReader, Write},
-    path::Path,
+    path::{Path, PathBuf},
     process::Command,
     time::{SystemTime, UNIX_EPOCH},
 };
@@ -266,9 +267,68 @@ impl Cheatcode for getDeployedCodeCall {
     }
 }
 
+/// Returns the path to the json artifact depending on the input
+fn get_artifact_path(state: &Cheatcodes, path: &str) -> Result<PathBuf> {
+    if path.ends_with(".json") {
+        Ok(PathBuf::from(path))
+    } else {
+        let mut parts = path.split(':');
+        let file = PathBuf::from(parts.next().unwrap());
+        let contract_name = parts.next();
+        let version = parts.next();
+
+        let version = if let Some(version) = version {
+            Some(Version::parse(version).map_err(|_| fmt_err!("Error parsing version"))?)
+        } else {
+            None
+        };
+
+        // Use available artifacts list if available
+        if let Some(available_ids) = &state.config.available_artifacts {
+            let mut artifact = None;
+
+            for id in available_ids.iter() {
+                // name might be in the form of "Counter.0.8.23"
+                let id_name = id.name.split(".").next().unwrap();
+                let id_source_path =
+                    id.source.strip_prefix(&state.config.root).unwrap_or(&id.source);
+
+                if file != id_source_path {
+                    continue;
+                }
+                if let Some(name) = contract_name {
+                    if id_name != name {
+                        continue;
+                    }
+                }
+                if let Some(ref version) = version {
+                    if id.version.minor != version.minor ||
+                        id.version.major != version.major ||
+                        id.version.patch != version.patch
+                    {
+                        continue;
+                    }
+                }
+                artifact = Some(id);
+            }
+
+            let artifact = artifact.ok_or_else(|| fmt_err!("No matching artifact found"))?;
+            Ok(artifact.path.clone())
+        } else {
+            let file = file.to_string_lossy();
+            let contract_name = if let Some(contract_name) = contract_name {
+                contract_name.to_owned()
+            } else {
+                file.replace(".sol", "")
+            };
+            Ok(state.config.paths.artifacts.join(format!("{file}/{contract_name}.json")))
+        }
+    }
+}
+
 /// Reads the bytecode object(s) from the matching artifact
 fn read_bytecode(state: &Cheatcodes, path: &str) -> Result<ContractObject> {
-    let path = get_artifact_path(&state.config.paths, path);
+    let path = get_artifact_path(state, path)?;
     let path = state.config.ensure_path_allowed(path, FsAccessKind::Read)?;
     let data = fs::read_to_string(path)?;
     serde_json::from_str::<ContractObject>(&data).map_err(Into::into)

crates/chisel/src/executor.rs

@@ -307,6 +307,7 @@ impl SessionSource {
                         &self.config.foundry_config,
                         self.config.evm_opts.clone(),
                         None,
+                        None,
                     )
                     .into(),
                 )

crates/common/src/contracts.rs

@@ -5,13 +5,12 @@ use alloy_primitives::{hex, Address, Selector, B256};
 use eyre::Result;
 use foundry_compilers::{
     artifacts::{CompactContractBytecode, ContractBytecodeSome},
-    ArtifactId, ProjectPathsConfig,
+    ArtifactId,
 };
 use std::{
     collections::BTreeMap,
     fmt,
     ops::{Deref, DerefMut},
-    path::PathBuf,
 };
 
 type ArtifactWithContractRef<'a> = (&'a ArtifactId, &'a (JsonAbi, Vec<u8>));
@@ -170,19 +169,6 @@ pub fn get_file_name(id: &str) -> &str {
     id.split(':').next().unwrap_or(id)
 }
 
-/// Returns the path to the json artifact depending on the input
-pub fn get_artifact_path(paths: &ProjectPathsConfig, path: &str) -> PathBuf {
-    if path.ends_with(".json") {
-        PathBuf::from(path)
-    } else {
-        let parts: Vec<&str> = path.split(':').collect();
-        let file = parts[0];
-        let contract_name =
-            if parts.len() == 1 { parts[0].replace(".sol", "") } else { parts[1].to_string() };
-        paths.artifacts.join(format!("{file}/{contract_name}.json"))
-    }
-}
-
 /// Helper function to convert CompactContractBytecode ~> ContractBytecodeSome
 pub fn compact_to_contract(contract: CompactContractBytecode) -> Result<ContractBytecodeSome> {
     Ok(ContractBytecodeSome {

crates/forge/bin/cmd/coverage.rs

@@ -293,14 +293,21 @@ impl CoverageArgs {
     ) -> Result<()> {
         let root = project.paths.root;
 
+        let artifact_ids = output.artifact_ids().map(|(id, _)| id).collect();
+
         // Build the contract runner
         let env = evm_opts.evm_env().await?;
         let mut runner = MultiContractRunnerBuilder::default()
             .initial_balance(evm_opts.initial_balance)
             .evm_spec(config.evm_spec_id())
             .sender(evm_opts.sender)
             .with_fork(evm_opts.get_fork(&config, env.clone()))
-            .with_cheats_config(CheatsConfig::new(&config, evm_opts.clone(), None))
+            .with_cheats_config(CheatsConfig::new(
+                &config,
+                evm_opts.clone(),
+                Some(artifact_ids),
+                None,
+            ))
             .with_test_options(TestOptions {
                 fuzz: config.fuzz,
                 invariant: config.invariant,

crates/forge/bin/cmd/test/mod.rs

@@ -254,13 +254,20 @@ impl TestArgs {
         // Clone the output only if we actually need it later for the debugger.
         let output_clone = should_debug.then(|| output.clone());
 
+        let artifact_ids = output.artifact_ids().map(|(id, _)| id).collect();
+
         let runner = MultiContractRunnerBuilder::default()
             .set_debug(should_debug)
             .initial_balance(evm_opts.initial_balance)
             .evm_spec(config.evm_spec_id())
             .sender(evm_opts.sender)
             .with_fork(evm_opts.get_fork(&config, env.clone()))
-            .with_cheats_config(CheatsConfig::new(&config, evm_opts.clone(), None))
+            .with_cheats_config(CheatsConfig::new(
+                &config,
+                evm_opts.clone(),
+                Some(artifact_ids),
+                None,
+            ))
             .with_test_options(test_options)
             .enable_isolation(evm_opts.isolate)
             .build(project_root, output, env, evm_opts)?;

crates/forge/tests/it/config.rs

@@ -142,8 +142,9 @@ pub fn runner_with_config(mut config: Config) -> MultiContractRunner {
     let opts = &*EVM_OPTS;
     let env = opts.local_evm_env();
     let output = COMPILED.clone();
+    let artifact_ids = output.artifact_ids().map(|(id, _)| id).collect();
     base_runner()
-        .with_cheats_config(CheatsConfig::new(&config, opts.clone(), None))
+        .with_cheats_config(CheatsConfig::new(&config, opts.clone(), Some(artifact_ids), None))
         .sender(config.sender)
         .build(root, output, env, opts.clone())
         .unwrap()

crates/script/src/build.rs

@@ -26,6 +26,9 @@ pub struct BuildData {
     pub target: ArtifactId,
     /// Source files of the contracts. Used by debugger.
     pub sources: ContractSources,
+    /// Artifact ids of the contracts. Passed to cheatcodes to enable usage of
+    /// `vm.getDeployedCode`.
+    pub artifact_ids: Vec<ArtifactId>,
 }
 
 impl BuildData {
@@ -181,6 +184,8 @@ impl PreprocessedState {
 
         let mut target_id: Option<ArtifactId> = None;
 
+        let artifact_ids = output.artifact_ids().map(|(id, _)| id).collect();
+
         // Find target artfifact id by name and path in compilation artifacts.
         for (id, contract) in output.artifact_ids().filter(|(id, _)| id.source == target_path) {
             if let Some(name) = &target_name {
@@ -220,7 +225,7 @@ impl PreprocessedState {
             args,
             script_config,
             script_wallets,
-            build_data: BuildData { linker, target, sources },
+            build_data: BuildData { linker, artifact_ids, target, sources },
         })
     }
 }

crates/script/src/execute.rs

@@ -98,7 +98,11 @@ impl PreExecutionState {
     pub async fn execute(mut self) -> Result<ExecutedState> {
         let mut runner = self
             .script_config
-            .get_runner_with_cheatcodes(self.script_wallets.clone(), self.args.debug)
+            .get_runner_with_cheatcodes(
+                self.build_data.build_data.artifact_ids.clone(),
+                self.script_wallets.clone(),
+                self.args.debug,
+            )
             .await?;
         let mut result = self.execute_with_runner(&mut runner).await?;
 

crates/script/src/lib.rs

@@ -541,15 +541,16 @@ impl ScriptConfig {
 
     async fn get_runner_with_cheatcodes(
         &mut self,
+        artifact_ids: Vec<ArtifactId>,
         script_wallets: ScriptWallets,
         debug: bool,
     ) -> Result<ScriptRunner> {
-        self._get_runner(Some(script_wallets), debug).await
+        self._get_runner(Some((artifact_ids, script_wallets)), debug).await
     }
 
     async fn _get_runner(
         &mut self,
-        script_wallets: Option<ScriptWallets>,
+        cheats_data: Option<(Vec<ArtifactId>, ScriptWallets)>,
         debug: bool,
     ) -> Result<ScriptRunner> {
         trace!("preparing script runner");
@@ -578,14 +579,15 @@ impl ScriptConfig {
             .spec(self.config.evm_spec_id())
             .gas_limit(self.evm_opts.gas_limit());
 
-        if let Some(script_wallets) = script_wallets {
+        if let Some((artifact_ids, script_wallets)) = cheats_data {
             builder = builder.inspectors(|stack| {
                 stack
                     .debug(debug)
                     .cheatcodes(
                         CheatsConfig::new(
                             &self.config,
                             self.evm_opts.clone(),
+                            Some(artifact_ids),
                             Some(script_wallets),
                         )
                         .into(),

testdata/cheats/GetCode.t.sol

@@ -4,6 +4,8 @@ pragma solidity 0.8.18;
 import "ds-test/test.sol";
 import "./Vm.sol";
 
+contract TestContract {}
+
 contract GetCodeTest is DSTest {
     Vm constant vm = Vm(HEVM_ADDRESS);
 
@@ -70,4 +72,12 @@ contract GetCodeTest is DSTest {
     function testFailGetUnlinked() public {
         vm.getCode("UnlinkedContract.sol");
     }
+
+    function testWithVersion() public {
+        bytes memory code = vm.getCode("cheats/GetCode.t.sol:TestContract:0.8.18");
+        assertEq(type(TestContract).creationCode, code);
+
+        vm._expectCheatcodeRevert("No matching artifact found");
+        vm.getCode("cheats/GetCode.t.sol:TestContract:0.8.19");
+    }
 }

testdata/cheats/GetDeployedCode.t.sol

@@ -4,6 +4,8 @@ pragma solidity 0.8.18;
 import "ds-test/test.sol";
 import "./Vm.sol";
 
+contract TestContract {}
+
 contract GetDeployedCodeTest is DSTest {
     Vm constant vm = Vm(HEVM_ADDRESS);
 
@@ -36,6 +38,16 @@ contract GetDeployedCodeTest is DSTest {
         emit Payload(address(this), address(0), "hello");
         over.emitPayload(address(0), "hello");
     }
+
+    function testWithVersion() public {
+        TestContract test = new TestContract();
+        bytes memory code = vm.getDeployedCode("cheats/GetDeployedCode.t.sol:TestContract:0.8.18");
+
+        assertEq(address(test).code, code);
+
+        vm._expectCheatcodeRevert("No matching artifact found");
+        vm.getDeployedCode("cheats/GetDeployedCode.t.sol:TestContract:0.8.19");
+    }
 }
 
 interface Override {