fix: smarter getCode validation

crates/cheatcodes/src/config.rs

@@ -8,6 +8,7 @@ use foundry_config::{
     ResolvedRpcEndpoints,
 };
 use foundry_evm_core::opts::EvmOpts;
+use semver::Version;
 use std::{
     collections::HashMap,
     path::{Path, PathBuf},
@@ -47,6 +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<Vec<ArtifactId>>,
+    /// Version of the script/test contract which is currently running.
+    pub running_version: Option<Version>,
 }
 
 impl CheatsConfig {
@@ -56,6 +59,7 @@ impl CheatsConfig {
         evm_opts: EvmOpts,
         available_artifacts: Option<Vec<ArtifactId>>,
         script_wallets: Option<ScriptWallets>,
+        running_version: Option<Version>,
     ) -> Self {
         let mut allowed_paths = vec![config.__root.0.clone()];
         allowed_paths.extend(config.libs.clone());
@@ -82,6 +86,7 @@ impl CheatsConfig {
             labels: config.labels.clone(),
             script_wallets,
             available_artifacts,
+            running_version,
         }
     }
 
@@ -200,6 +205,7 @@ impl Default for CheatsConfig {
             labels: Default::default(),
             script_wallets: None,
             available_artifacts: Default::default(),
+            running_version: Default::default(),
         }
     }
 }
@@ -215,6 +221,7 @@ mod tests {
             Default::default(),
             None,
             None,
+            None,
         )
     }
 

crates/cheatcodes/src/fs.rs

@@ -288,35 +288,53 @@ fn get_artifact_path(state: &Cheatcodes, path: &str) -> Result<PathBuf> {
 
         // 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();
-
-                if !id.source.ends_with(&file) {
-                    continue;
-                }
-                if let Some(name) = contract_name {
-                    if id_name != name {
-                        continue;
+            let filtered = available_ids
+                .iter()
+                .filter(|id| {
+                    // name might be in the form of "Counter.0.8.23"
+                    let id_name = id.name.split('.').next().unwrap();
+
+                    if !id.source.ends_with(&file) {
+                        return false;
                     }
-                }
-                if let Some(ref version) = version {
-                    if id.version.minor != version.minor ||
-                        id.version.major != version.major ||
-                        id.version.patch != version.patch
-                    {
-                        continue;
+                    if let Some(name) = contract_name {
+                        if id_name != name {
+                            return false;
+                        }
                     }
+                    if let Some(ref version) = version {
+                        if id.version.minor != version.minor ||
+                            id.version.major != version.major ||
+                            id.version.patch != version.patch
+                        {
+                            return false;
+                        }
+                    }
+                    true
+                })
+                .collect::<Vec<_>>();
+
+            let artifact = match filtered.len() {
+                0 => Err(fmt_err!("No matching artifact found")),
+                1 => Ok(filtered[0]),
+                _ => {
+                    // 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
+                                .into_iter()
+                                .filter(|id| id.version == *version)
+                                .collect::<Vec<_>>();
+
+                            (filtered.len() == 1).then_some(filtered[0])
+                        })
+                        .ok_or_else(|| fmt_err!("Multiple matching artifacts found"))
                 }
-                if artifact.is_some() {
-                    return Err(fmt_err!("Multiple matching artifacts found"));
-                }
-                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();

crates/chisel/src/executor.rs

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

crates/forge/bin/cmd/coverage.rs

@@ -317,6 +317,7 @@ impl CoverageArgs {
                 evm_opts.clone(),
                 Some(artifact_ids),
                 None,
+                None,
             ))
             .with_test_options(TestOptions {
                 fuzz: config.fuzz,

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

@@ -287,6 +287,7 @@ impl TestArgs {
                 evm_opts.clone(),
                 Some(artifact_ids),
                 None,
+                None, // populated separately for each test contract
             ))
             .with_test_options(test_options)
             .enable_isolation(evm_opts.isolate)

crates/forge/src/multi_runner.rs

@@ -9,13 +9,8 @@ use foundry_compilers::{
     artifacts::Libraries, contracts::ArtifactContracts, Artifact, ArtifactId, ProjectCompileOutput,
 };
 use foundry_evm::{
-    backend::Backend,
-    decode::RevertDecoder,
-    executors::{Executor, ExecutorBuilder},
-    fork::CreateFork,
-    inspectors::CheatsConfig,
-    opts::EvmOpts,
-    revm,
+    backend::Backend, decode::RevertDecoder, executors::ExecutorBuilder, fork::CreateFork,
+    inspectors::CheatsConfig, opts::EvmOpts, revm,
 };
 use foundry_linking::{LinkOutput, Linker};
 use rayon::prelude::*;
@@ -158,18 +153,6 @@ impl MultiContractRunner {
 
         // The DB backend that serves all the data.
         let db = Backend::spawn(self.fork.take());
-        let executor = ExecutorBuilder::new()
-            .inspectors(|stack| {
-                stack
-                    .cheatcodes(self.cheats_config.clone())
-                    .trace(self.evm_opts.verbosity >= 3 || self.debug)
-                    .debug(self.debug)
-                    .coverage(self.coverage)
-                    .enable_isolation(self.isolation)
-            })
-            .spec(self.evm_spec)
-            .gas_limit(self.evm_opts.gas_limit())
-            .build(self.env.clone(), db);
 
         let find_timer = Instant::now();
         let contracts = self.matching_contracts(filter).collect::<Vec<_>>();
@@ -182,30 +165,46 @@ impl MultiContractRunner {
         );
 
         contracts.par_iter().for_each_with(tx, |tx, &(id, contract)| {
-            let identifier = id.identifier();
-            let executor = executor.clone();
-            let result = self.run_tests(&identifier, contract, executor, filter);
-            let _ = tx.send((identifier, result));
+            let result = self.run_tests(id, contract, db.clone(), filter);
+            let _ = tx.send((id.identifier(), result));
         })
     }
 
     fn run_tests(
         &self,
-        name: &str,
+        artifact_id: &ArtifactId,
         contract: &TestContract,
-        executor: Executor,
+        db: Backend,
         filter: &dyn TestFilter,
     ) -> SuiteResult {
-        let mut span_name = name;
+        let identifier = artifact_id.identifier();
+        let mut span_name = identifier.as_str();
+
+        let mut cheats_config = self.cheats_config.as_ref().clone();
+        cheats_config.running_version = Some(artifact_id.version.clone());
+
+        let executor = ExecutorBuilder::new()
+            .inspectors(|stack| {
+                stack
+                    .cheatcodes(Arc::new(cheats_config))
+                    .trace(self.evm_opts.verbosity >= 3 || self.debug)
+                    .debug(self.debug)
+                    .coverage(self.coverage)
+                    .enable_isolation(self.isolation)
+            })
+            .spec(self.evm_spec)
+            .gas_limit(self.evm_opts.gas_limit())
+            .build(self.env.clone(), db);
+
         if !enabled!(tracing::Level::TRACE) {
-            span_name = get_contract_name(span_name);
+            span_name = get_contract_name(&identifier);
         }
         let _guard = info_span!("run_tests", name = span_name).entered();
 
         debug!("start executing all tests in contract");
 
         let runner = ContractRunner::new(
-            name,
+            &identifier,
             executor,
             contract,
             self.evm_opts.initial_balance,

crates/forge/tests/it/test_helpers.rs

@@ -206,7 +206,13 @@ impl ForgeTestData {
         let output = self.output.clone();
         let artifact_ids = output.artifact_ids().map(|(id, _)| id).collect();
         self.base_runner()
-            .with_cheats_config(CheatsConfig::new(&config, opts.clone(), Some(artifact_ids), None))
+            .with_cheats_config(CheatsConfig::new(
+                &config,
+                opts.clone(),
+                Some(artifact_ids),
+                None,
+                None,
+            ))
             .sender(config.sender)
             .with_test_options(self.test_opts.clone())
             .build(root, output, env, opts.clone())

crates/script/src/execute.rs

@@ -102,6 +102,7 @@ impl PreExecutionState {
                 self.build_data.build_data.artifact_ids.clone(),
                 self.script_wallets.clone(),
                 self.args.debug,
+                self.build_data.build_data.target.clone(),
             )
             .await?;
         let mut result = self.execute_with_runner(&mut runner).await?;

crates/script/src/lib.rs

@@ -559,13 +559,14 @@ impl ScriptConfig {
         artifact_ids: Vec<ArtifactId>,
         script_wallets: ScriptWallets,
         debug: bool,
+        target: ArtifactId,
     ) -> Result<ScriptRunner> {
-        self._get_runner(Some((artifact_ids, script_wallets)), debug).await
+        self._get_runner(Some((artifact_ids, script_wallets, target)), debug).await
     }
 
     async fn _get_runner(
         &mut self,
-        cheats_data: Option<(Vec<ArtifactId>, ScriptWallets)>,
+        cheats_data: Option<(Vec<ArtifactId>, ScriptWallets, ArtifactId)>,
         debug: bool,
     ) -> Result<ScriptRunner> {
         trace!("preparing script runner");
@@ -594,7 +595,7 @@ impl ScriptConfig {
             .spec(self.config.evm_spec_id())
             .gas_limit(self.evm_opts.gas_limit());
 
-        if let Some((artifact_ids, script_wallets)) = cheats_data {
+        if let Some((artifact_ids, script_wallets, target)) = cheats_data {
             builder = builder.inspectors(|stack| {
                 stack
                     .debug(debug)
@@ -604,6 +605,7 @@ impl ScriptConfig {
                             self.evm_opts.clone(),
                             Some(artifact_ids),
                             Some(script_wallets),
+                            Some(target.version),
                         )
                         .into(),
                     )