fix: better artifacts management for getCode

crates/cheatcodes/src/config.rs

@@ -1,8 +1,8 @@
 use super::Result;
 use crate::{script::ScriptWallets, Vm::Rpc};
 use alloy_primitives::Address;
-use foundry_common::fs::normalize_path;
-use foundry_compilers::{utils::canonicalize, ArtifactId, ProjectPathsConfig};
+use foundry_common::{fs::normalize_path, ContractsByArtifact};
+use foundry_compilers::{utils::canonicalize, ProjectPathsConfig};
 use foundry_config::{
     cache::StorageCachingConfig, fs_permissions::FsAccessKind, Config, FsPermissions,
     ResolvedRpcEndpoints,
@@ -12,6 +12,7 @@ use semver::Version;
 use std::{
     collections::HashMap,
     path::{Path, PathBuf},
+    sync::Arc,
     time::Duration,
 };
 
@@ -47,7 +48,7 @@ pub struct CheatsConfig {
     /// 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>>,
+    pub available_artifacts: Option<Arc<ContractsByArtifact>>,
     /// Version of the script/test contract which is currently running.
     pub running_version: Option<Version>,
 }
@@ -57,7 +58,7 @@ impl CheatsConfig {
     pub fn new(
         config: &Config,
         evm_opts: EvmOpts,
-        available_artifacts: Option<Vec<ArtifactId>>,
+        available_artifacts: Option<Arc<ContractsByArtifact>>,
         script_wallets: Option<ScriptWallets>,
         running_version: Option<Version>,
     ) -> Self {

crates/cheatcodes/src/fs.rs

@@ -4,7 +4,7 @@ use super::string::parse;
 use crate::{Cheatcode, Cheatcodes, Result, Vm::*};
 use alloy_dyn_abi::DynSolType;
 use alloy_json_abi::ContractObject;
-use alloy_primitives::U256;
+use alloy_primitives::{Bytes, U256};
 use alloy_sol_types::SolValue;
 use dialoguer::{Input, Password};
 use foundry_common::fs;
@@ -251,24 +251,14 @@ impl Cheatcode for writeLineCall {
 impl Cheatcode for getCodeCall {
     fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self { artifactPath: path } = self;
-        let object = read_bytecode(state, path)?;
-        if let Some(bin) = object.bytecode {
-            Ok(bin.abi_encode())
-        } else {
-            Err(fmt_err!("No bytecode for contract. Is it abstract or unlinked?"))
-        }
+        Ok(get_artifact_code(state, path, false)?.abi_encode())
     }
 }
 
 impl Cheatcode for getDeployedCodeCall {
     fn apply(&self, state: &mut Cheatcodes) -> Result {
         let Self { artifactPath: path } = self;
-        let object = read_bytecode(state, path)?;
-        if let Some(bin) = object.deployed_bytecode {
-            Ok(bin.abi_encode())
-        } else {
-            Err(fmt_err!("No deployed bytecode for contract. Is it abstract or unlinked?"))
-        }
+        Ok(get_artifact_code(state, path, true)?.abi_encode())
     }
 }
 
@@ -282,9 +272,9 @@ impl Cheatcode for getDeployedCodeCall {
 /// - `path/to/contract.sol:0.8.23`
 /// - `ContractName`
 /// - `ContractName:0.8.23`
-fn get_artifact_path(state: &Cheatcodes, path: &str) -> Result<PathBuf> {
-    if path.ends_with(".json") {
-        Ok(PathBuf::from(path))
+fn get_artifact_code(state: &Cheatcodes, path: &str, deployed: bool) -> Result<Bytes> {
+    let path = if path.ends_with(".json") {
+        PathBuf::from(path)
     } else {
         let mut parts = path.split(':');
 
@@ -314,11 +304,11 @@ fn get_artifact_path(state: &Cheatcodes, path: &str) -> Result<PathBuf> {
             None
         };
 
-        // Use available artifacts list if available
-        if let Some(available_ids) = &state.config.available_artifacts {
-            let filtered = available_ids
+        // Use available artifacts list if present
+        if let Some(artifacts) = &state.config.available_artifacts {
+            let filtered = artifacts
                 .iter()
-                .filter(|id| {
+                .filter(|(id, _)| {
                     // name might be in the form of "Counter.0.8.23"
                     let id_name = id.name.split('.').next().unwrap();
 
@@ -356,7 +346,7 @@ fn get_artifact_path(state: &Cheatcodes, path: &str) -> Result<PathBuf> {
                         .and_then(|version| {
                             let filtered = filtered
                                 .into_iter()
-                                .filter(|id| id.version == *version)
+                                .filter(|(id, _)| id.version == *version)
                                 .collect::<Vec<_>>();
 
                             (filtered.len() == 1).then_some(filtered[0])
@@ -365,31 +355,33 @@ fn get_artifact_path(state: &Cheatcodes, path: &str) -> Result<PathBuf> {
                 }
             }?;
 
-            Ok(artifact.path.clone())
+            if deployed {
+                return Ok(artifact.1.deployed_bytecode.clone())
+            } else {
+                return Ok(artifact.1.bytecode.clone())
+            }
         } else {
-            let path_in_artifacts =
-                match (file.map(|f| f.to_string_lossy().to_string()), contract_name) {
-                    (Some(file), Some(contract_name)) => Ok(format!("{file}/{contract_name}.json")),
-                    (None, Some(contract_name)) => {
-                        Ok(format!("{contract_name}.sol/{contract_name}.json"))
-                    }
-                    (Some(file), None) => {
-                        let name = file.replace(".sol", "");
-                        Ok(format!("{file}/{name}.json"))
-                    }
-                    _ => Err(fmt_err!("Invalid artifact path")),
-                }?;
-            Ok(state.config.paths.artifacts.join(path_in_artifacts))
+            match (file.map(|f| f.to_string_lossy().to_string()), contract_name) {
+                (Some(file), Some(contract_name)) => {
+                    PathBuf::from(format!("{file}/{contract_name}.json"))
+                }
+                (None, Some(contract_name)) => {
+                    PathBuf::from(format!("{contract_name}.sol/{contract_name}.json"))
+                }
+                (Some(file), None) => {
+                    let name = file.replace(".sol", "");
+                    PathBuf::from(format!("{file}/{name}.json"))
+                }
+                _ => return Err(fmt_err!("Invalid artifact path")),
+            }
         }
-    }
-}
+    };
 
-/// Reads the bytecode object(s) from the matching artifact
-fn read_bytecode(state: &Cheatcodes, path: &str) -> Result<ContractObject> {
-    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)
+    let artifact = serde_json::from_str::<ContractObject>(&data)?;
+    let maybe_bytecode = if deployed { artifact.deployed_bytecode } else { artifact.bytecode };
+    maybe_bytecode.ok_or_else(|| fmt_err!("No bytecode for contract. Is it abstract or unlinked?"))
 }
 
 impl Cheatcode for ffiCall {

crates/common/src/contracts.rs

@@ -32,6 +32,26 @@ type ArtifactWithContractRef<'a> = (&'a ArtifactId, &'a ContractData);
 pub struct ContractsByArtifact(pub BTreeMap<ArtifactId, ContractData>);
 
 impl ContractsByArtifact {
+    /// Creates a new instance by collecting all artifacts with present bytecode from an iterator.
+    ///
+    /// It is recommended to use this method with an output of
+    /// [foundry_linking::Linker::get_linked_artifacts].
+    pub fn new(artifacts: impl IntoIterator<Item = (ArtifactId, CompactContractBytecode)>) -> Self {
+        Self(
+            artifacts
+                .into_iter()
+                .filter_map(|(id, artifact)| {
+                    let name = id.name.clone();
+                    let bytecode = artifact.bytecode.and_then(|b| b.into_bytes())?;
+                    let deployed_bytecode =
+                        artifact.deployed_bytecode.and_then(|b| b.into_bytes())?;
+                    let abi = artifact.abi?;
+
+                    Some((id, ContractData { name, abi, bytecode, deployed_bytecode }))
+                })
+                .collect(),
+        )
+    }
     /// Finds a contract which has a similar bytecode as `code`.
     pub fn find_by_creation_code(&self, code: &[u8]) -> Option<ArtifactWithContractRef> {
         self.iter()

crates/evm/evm/src/executors/invariant/error.rs

@@ -149,7 +149,7 @@ impl FailedInvariantCaseData {
     pub fn replay(
         &self,
         mut executor: Executor,
-        known_contracts: Option<&ContractsByArtifact>,
+        known_contracts: &ContractsByArtifact,
         mut ided_contracts: ContractsByAddress,
         logs: &mut Vec<Log>,
         traces: &mut Traces,

crates/evm/evm/src/executors/invariant/funcs.rs

@@ -71,7 +71,7 @@ pub fn assert_invariants(
 pub fn replay_run(
     invariant_contract: &InvariantContract<'_>,
     mut executor: Executor,
-    known_contracts: Option<&ContractsByArtifact>,
+    known_contracts: &ContractsByArtifact,
     mut ided_contracts: ContractsByAddress,
     logs: &mut Vec<Log>,
     traces: &mut Traces,

crates/evm/traces/src/lib.rs

@@ -12,7 +12,7 @@ use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact};
 use foundry_evm_core::constants::CHEATCODE_ADDRESS;
 use futures::{future::BoxFuture, FutureExt};
 use serde::{Deserialize, Serialize};
-use std::{collections::BTreeMap, fmt::Write};
+use std::fmt::Write;
 use yansi::{Color, Paint};
 
 /// Call trace address identifiers.
@@ -293,12 +293,8 @@ fn trace_color(trace: &CallTrace) -> Color {
 }
 
 /// Given a list of traces and artifacts, it returns a map connecting address to abi
-pub fn load_contracts(
-    traces: Traces,
-    known_contracts: Option<&ContractsByArtifact>,
-) -> ContractsByAddress {
-    let Some(contracts) = known_contracts else { return BTreeMap::new() };
-    let mut local_identifier = LocalTraceIdentifier::new(contracts);
+pub fn load_contracts(traces: Traces, known_contracts: &ContractsByArtifact) -> ContractsByAddress {
+    let mut local_identifier = LocalTraceIdentifier::new(known_contracts);
     let mut decoder = CallTraceDecoderBuilder::new().build();
     for (_, trace) in &traces {
         decoder.identify(trace, &mut local_identifier);
@@ -308,7 +304,7 @@ pub fn load_contracts(
         .contracts
         .iter()
         .filter_map(|(addr, name)| {
-            if let Ok(Some((_, contract))) = contracts.find_by_name_or_identifier(name) {
+            if let Ok(Some((_, contract))) = known_contracts.find_by_name_or_identifier(name) {
                 return Some((*addr, (name.clone(), contract.abi.clone())));
             }
             None

crates/forge/bin/cmd/coverage.rs

@@ -7,7 +7,6 @@ use forge::{
         analysis::SourceAnalyzer, anchors::find_anchors, BytecodeReporter, ContractId,
         CoverageReport, CoverageReporter, DebugReporter, LcovReporter, SummaryReporter,
     },
-    inspectors::CheatsConfig,
     opts::EvmOpts,
     result::SuiteResult,
     revm::primitives::SpecId,
@@ -313,22 +312,13 @@ 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()
+        let mut runner = MultiContractRunnerBuilder::new(config.clone())
             .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(),
-                Some(artifact_ids),
-                None,
-                None,
-            ))
             .with_test_options(TestOptions {
                 fuzz: config.fuzz,
                 invariant: config.invariant,
@@ -338,31 +328,32 @@ impl CoverageArgs {
             .build(&root, output, env, evm_opts)?;
 
         // Run tests
-        let known_contracts = runner.known_contracts.clone();
         let filter = self.filter;
         let (tx, rx) = channel::<(String, SuiteResult)>();
         let handle = tokio::task::spawn_blocking(move || runner.test(&filter, tx));
 
         // Add hit data to the coverage report
-        let data = rx
-            .into_iter()
-            .flat_map(|(_, suite)| suite.test_results.into_values())
-            .filter_map(|mut result| result.coverage.take())
-            .flat_map(|hit_maps| {
-                hit_maps.0.into_values().filter_map(|map| {
+        let data = rx.into_iter().flat_map(|(_, suite)| {
+            let mut hits = Vec::new();
+            for (_, mut result) in suite.test_results {
+                let Some(hit_maps) = result.coverage.take() else { continue };
+
+                for map in hit_maps.0.into_values() {
                     if let Some((id, _)) =
-                        known_contracts.find_by_deployed_code(map.bytecode.as_ref())
+                        suite.known_contracts.find_by_deployed_code(map.bytecode.as_ref())
                     {
-                        Some((id, map, true))
+                        hits.push((id.clone(), map, true));
                     } else if let Some((id, _)) =
-                        known_contracts.find_by_creation_code(map.bytecode.as_ref())
+                        suite.known_contracts.find_by_creation_code(map.bytecode.as_ref())
                     {
-                        Some((id, map, false))
-                    } else {
-                        None
+                        hits.push((id.clone(), map, false));
                     }
-                })
-            });
+                }
+            }
+
+            hits
+        });
+
         for (artifact_id, hits, is_deployed_code) in data {
             // TODO: Note down failing tests
             if let Some(source_id) = report.get_source_id(