feat: filter by profile in vm.getCode (#9714)

feat: filter by profile in getCode

crates/cheatcodes/src/config.rs

@@ -2,13 +2,12 @@ use super::Result;
 use crate::Vm::Rpc;
 use alloy_primitives::{map::AddressHashMap, U256};
 use foundry_common::{fs::normalize_path, ContractsByArtifact};
-use foundry_compilers::{utils::canonicalize, ProjectPathsConfig};
+use foundry_compilers::{utils::canonicalize, ArtifactId, ProjectPathsConfig};
 use foundry_config::{
     cache::StorageCachingConfig, fs_permissions::FsAccessKind, Config, FsPermissions,
     ResolvedRpcEndpoint, ResolvedRpcEndpoints, RpcEndpoint, RpcEndpointUrl,
 };
 use foundry_evm_core::opts::EvmOpts;
-use semver::Version;
 use std::{
     path::{Path, PathBuf},
     time::Duration,
@@ -49,10 +48,8 @@ pub struct CheatsConfig {
     /// If Some, `vm.getDeployedCode` invocations are validated to be in scope of this list.
     /// If None, no validation is performed.
     pub available_artifacts: Option<ContractsByArtifact>,
-    /// Name of the script/test contract which is currently running.
-    pub running_contract: Option<String>,
-    /// Version of the script/test contract which is currently running.
-    pub running_version: Option<Version>,
+    /// Currently running artifact.
+    pub running_artifact: Option<ArtifactId>,
     /// Whether to enable legacy (non-reverting) assertions.
     pub assertions_revert: bool,
     /// Optional seed for the RNG algorithm.
@@ -65,8 +62,7 @@ impl CheatsConfig {
         config: &Config,
         evm_opts: EvmOpts,
         available_artifacts: Option<ContractsByArtifact>,
-        running_contract: Option<String>,
-        running_version: Option<Version>,
+        running_artifact: Option<ArtifactId>,
     ) -> Self {
         let mut allowed_paths = vec![config.root.clone()];
         allowed_paths.extend(config.libs.iter().cloned());
@@ -94,22 +90,15 @@ impl CheatsConfig {
             evm_opts,
             labels: config.labels.clone(),
             available_artifacts,
-            running_contract,
-            running_version,
+            running_artifact,
             assertions_revert: config.assertions_revert,
             seed: config.fuzz.seed,
         }
     }
 
     /// Returns a new `CheatsConfig` configured with the given `Config` and `EvmOpts`.
     pub fn clone_with(&self, config: &Config, evm_opts: EvmOpts) -> Self {
-        Self::new(
-            config,
-            evm_opts,
-            self.available_artifacts.clone(),
-            self.running_contract.clone(),
-            self.running_version.clone(),
-        )
+        Self::new(config, evm_opts, self.available_artifacts.clone(), self.running_artifact.clone())
     }
 
     /// Attempts to canonicalize (see [std::fs::canonicalize]) the path.
@@ -230,8 +219,7 @@ impl Default for CheatsConfig {
             evm_opts: Default::default(),
             labels: Default::default(),
             available_artifacts: Default::default(),
-            running_contract: Default::default(),
-            running_version: Default::default(),
+            running_artifact: Default::default(),
             assertions_revert: true,
             seed: None,
         }
@@ -249,7 +237,6 @@ mod tests {
             Default::default(),
             None,
             None,
-            None,
         )
     }
 

crates/cheatcodes/src/evm.rs

@@ -1038,7 +1038,7 @@ fn derive_snapshot_name(
     name: Option<String>,
 ) -> (String, String) {
     let group = group.unwrap_or_else(|| {
-        ccx.state.config.running_contract.clone().expect("expected running contract")
+        ccx.state.config.running_artifact.clone().expect("expected running contract").name
     });
     let name = name.unwrap_or_else(|| "default".to_string());
     (group, name)

crates/cheatcodes/src/fs.rs

@@ -413,19 +413,31 @@ fn get_artifact_code(state: &Cheatcodes, path: &str, deployed: bool) -> Result<B
 
             let artifact = match &filtered[..] {
                 [] => Err(fmt_err!("no matching artifact found")),
-                [artifact] => Ok(artifact),
+                [artifact] => Ok(*artifact),
                 filtered => {
+                    let mut filtered = filtered.to_vec();
                     // If we know the current script/test contract solc version, try to filter by it
                     state
                         .config
-                        .running_version
+                        .running_artifact
                         .as_ref()
-                        .and_then(|version| {
-                            let filtered = filtered
-                                .iter()
-                                .filter(|(id, _)| id.version == *version)
-                                .collect::<Vec<_>>();
-                            (filtered.len() == 1).then(|| filtered[0])
+                        .and_then(|running| {
+                            // Firstly filter by version
+                            filtered.retain(|(id, _)| id.version == running.version);
+
+                            // Return artifact if only one matched
+                            if filtered.len() == 1 {
+                                return Some(filtered[0])
+                            }
+
+                            // Try filtering by profile as well
+                            filtered.retain(|(id, _)| id.profile == running.profile);
+
+                            if filtered.len() == 1 {
+                                Some(filtered[0])
+                            } else {
+                                None
+                            }
                         })
                         .ok_or_else(|| fmt_err!("multiple matching artifacts found"))
                 }

crates/chisel/src/executor.rs

@@ -335,7 +335,6 @@ impl SessionSource {
                         self.config.evm_opts.clone(),
                         None,
                         None,
-                        Some(self.solc.version.clone()),
                     )
                     .into(),
                 )

crates/forge/src/multi_runner.rs

@@ -343,8 +343,7 @@ impl TestRunnerConfig {
             &self.config,
             self.evm_opts.clone(),
             Some(known_contracts),
-            Some(artifact_id.name.clone()),
-            Some(artifact_id.version.clone()),
+            Some(artifact_id.clone()),
         ));
         ExecutorBuilder::new()
             .inspectors(|stack| {

crates/script/src/lib.rs

@@ -629,8 +629,7 @@ impl ScriptConfig {
                             &self.config,
                             self.evm_opts.clone(),
                             Some(known_contracts),
-                            Some(target.name),
-                            Some(target.version),
+                            Some(target),
                         )
                         .into(),
                     )