chore: Arc ContractsByArtifact internally (#8026)

* chore: Arc ContractsByArtifact internally

* perf: clear from test suite result

* chore: clear only when not coverage

crates/cheatcodes/src/config.rs

@@ -12,7 +12,6 @@ use semver::Version;
 use std::{
     collections::HashMap,
     path::{Path, PathBuf},
-    sync::Arc,
     time::Duration,
 };
 
@@ -50,7 +49,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<Arc<ContractsByArtifact>>,
+    pub available_artifacts: Option<ContractsByArtifact>,
     /// Version of the script/test contract which is currently running.
     pub running_version: Option<Version>,
 }
@@ -60,7 +59,7 @@ impl CheatsConfig {
     pub fn new(
         config: &Config,
         evm_opts: EvmOpts,
-        available_artifacts: Option<Arc<ContractsByArtifact>>,
+        available_artifacts: Option<ContractsByArtifact>,
         script_wallets: Option<ScriptWallets>,
         running_version: Option<Version>,
     ) -> Self {

crates/common/src/contracts.rs

@@ -10,11 +10,7 @@ use foundry_compilers::{
     },
     ArtifactId,
 };
-use std::{
-    collections::BTreeMap,
-    ops::{Deref, DerefMut},
-    str::FromStr,
-};
+use std::{collections::BTreeMap, ops::Deref, str::FromStr, sync::Arc};
 
 /// Libraries' runtime code always starts with the following instruction:
 /// `PUSH20 0x0000000000000000000000000000000000000000`
@@ -60,7 +56,7 @@ impl From<CompactDeployedBytecode> for BytecodeData {
 }
 
 /// Container for commonly used contract data.
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct ContractData {
     /// Contract name.
     pub name: String,
@@ -88,34 +84,36 @@ type ArtifactWithContractRef<'a> = (&'a ArtifactId, &'a ContractData);
 
 /// Wrapper type that maps an artifact to a contract ABI and bytecode.
 #[derive(Clone, Default, Debug)]
-pub struct ContractsByArtifact(pub BTreeMap<ArtifactId, ContractData>);
+pub struct ContractsByArtifact(Arc<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 CompactContractBytecode { abi, bytecode, deployed_bytecode } = artifact;
-
-                    Some((
-                        id,
-                        ContractData {
-                            name,
-                            abi: abi?,
-                            bytecode: bytecode.map(Into::into),
-                            deployed_bytecode: deployed_bytecode.map(Into::into),
-                        },
-                    ))
-                })
-                .collect(),
-        )
+        let map = artifacts
+            .into_iter()
+            .filter_map(|(id, artifact)| {
+                let name = id.name.clone();
+                let CompactContractBytecode { abi, bytecode, deployed_bytecode } = artifact;
+                Some((
+                    id,
+                    ContractData {
+                        name,
+                        abi: abi?,
+                        bytecode: bytecode.map(Into::into),
+                        deployed_bytecode: deployed_bytecode.map(Into::into),
+                    },
+                ))
+            })
+            .collect();
+        Self(Arc::new(map))
+    }
+
+    /// Clears all contracts.
+    pub fn clear(&mut self) {
+        *self = Self::default();
     }
 
     /// Finds a contract which has a similar bytecode as `code`.
@@ -276,12 +274,6 @@ impl Deref for ContractsByArtifact {
     }
 }
 
-impl DerefMut for ContractsByArtifact {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.0
-    }
-}
-
 /// Wrapper type that maps an address to a contract identifier and contract ABI.
 pub type ContractsByAddress = BTreeMap<Address, (String, JsonAbi)>;
 

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

@@ -392,12 +392,12 @@ impl TestArgs {
         let mut outcome = TestOutcome::empty(self.allow_failure);
 
         let mut any_test_failed = false;
-        for (contract_name, suite_result) in rx {
+        for (contract_name, mut suite_result) in rx {
             let tests = &suite_result.test_results;
 
             // Set up trace identifiers.
-            let known_contracts = suite_result.known_contracts.clone();
-            let mut identifier = TraceIdentifiers::new().with_local(&known_contracts);
+            let known_contracts = &suite_result.known_contracts;
+            let mut identifier = TraceIdentifiers::new().with_local(known_contracts);
 
             // Avoid using etherscan for gas report as we decode more traces and this will be
             // expensive.
@@ -407,7 +407,7 @@ impl TestArgs {
 
             // Build the trace decoder.
             let mut builder = CallTraceDecoderBuilder::new()
-                .with_known_contracts(&known_contracts)
+                .with_known_contracts(known_contracts)
                 .with_verbosity(verbosity);
             // Signatures are of no value for gas reports.
             if !self.gas_report {
@@ -453,10 +453,6 @@ impl TestArgs {
                 // processing the remaining tests and print the suite summary.
                 any_test_failed |= result.status == TestStatus::Failure;
 
-                if result.traces.is_empty() {
-                    continue;
-                }
-
                 // Clear the addresses and labels from previous runs.
                 decoder.clear_addresses();
                 decoder
@@ -524,6 +520,9 @@ impl TestArgs {
             // Print suite summary.
             shell::println(suite_result.summary())?;
 
+            // Free memory if it's not needed.
+            suite_result.clear_unneeded();
+
             // Add the suite result to the outcome.
             outcome.results.insert(contract_name, suite_result);
             outcome.last_run_decoder = Some(decoder);

crates/forge/src/multi_runner.rs

@@ -186,7 +186,7 @@ impl MultiContractRunner {
             self.output.artifact_ids().collect(),
         );
         let linked_contracts = linker.get_linked_artifacts(&contract.libraries).unwrap_or_default();
-        let known_contracts = Arc::new(ContractsByArtifact::new(linked_contracts));
+        let known_contracts = ContractsByArtifact::new(linked_contracts);
 
         let cheats_config = CheatsConfig::new(
             &self.config,

crates/forge/src/result.rs

@@ -1,5 +1,6 @@
 //! Test outcomes.
 
+use crate::gas_report::GasReport;
 use alloy_primitives::{Address, Log};
 use foundry_common::{
     evm::Breakpoints, get_contract_name, get_file_name, shell, ContractsByArtifact,
@@ -16,13 +17,10 @@ use serde::{Deserialize, Serialize};
 use std::{
     collections::{BTreeMap, HashMap},
     fmt::{self, Write},
-    sync::Arc,
     time::Duration,
 };
 use yansi::Paint;
 
-use crate::gas_report::GasReport;
-
 /// The aggregated result of a test run.
 #[derive(Clone, Debug)]
 pub struct TestOutcome {
@@ -201,8 +199,10 @@ pub struct SuiteResult {
     /// Libraries used to link test contract.
     pub libraries: Libraries,
     /// Contracts linked with correct libraries.
+    ///
+    /// This is cleared at the end of the test run if coverage is not enabled.
     #[serde(skip)]
-    pub known_contracts: Arc<ContractsByArtifact>,
+    pub known_contracts: ContractsByArtifact,
 }
 
 impl SuiteResult {
@@ -211,11 +211,18 @@ impl SuiteResult {
         test_results: BTreeMap<String, TestResult>,
         warnings: Vec<String>,
         libraries: Libraries,
-        known_contracts: Arc<ContractsByArtifact>,
+        known_contracts: ContractsByArtifact,
     ) -> Self {
         Self { duration, test_results, warnings, libraries, known_contracts }
     }
 
+    /// Frees memory that is not used for the final output.
+    pub fn clear_unneeded(&mut self) {
+        if !self.test_results.values().any(|r| r.coverage.is_some()) {
+            ContractsByArtifact::clear(&mut self.known_contracts);
+        }
+    }
+
     /// Returns an iterator over all individual succeeding tests and their names.
     pub fn successes(&self) -> impl Iterator<Item = (&String, &TestResult)> {
         self.tests().filter(|(_, t)| t.status == TestStatus::Success)

crates/forge/src/runner.rs

@@ -40,7 +40,6 @@ use std::{
     borrow::Cow,
     cmp::min,
     collections::{BTreeMap, HashMap},
-    sync::Arc,
     time::Instant,
 };
 
@@ -258,7 +257,7 @@ impl<'a> ContractRunner<'a> {
         mut self,
         filter: &dyn TestFilter,
         test_options: &TestOptions,
-        known_contracts: Arc<ContractsByArtifact>,
+        known_contracts: ContractsByArtifact,
         handle: &tokio::runtime::Handle,
     ) -> SuiteResult {
         info!("starting tests");

crates/forge/tests/cli/coverage.rs

@@ -70,11 +70,9 @@ contract AContractTest is DSTest {
     let lcov_data = std::fs::read_to_string(lcov_info).unwrap();
     // AContract.init must be hit at least once
     let re = Regex::new(r"FNDA:(\d+),AContract\.init").unwrap();
-    assert!(lcov_data.lines().any(|line| re.captures(line).map_or(false, |caps| caps
-        .get(1)
-        .unwrap()
-        .as_str()
-        .parse::<i32>()
-        .unwrap() >
-        0)));
+    let valid_line = |line| {
+        re.captures(line)
+            .map_or(false, |caps| caps.get(1).unwrap().as_str().parse::<i32>().unwrap() > 0)
+    };
+    assert!(lcov_data.lines().any(valid_line), "{lcov_data}");
 });

crates/script/src/build.rs

@@ -146,10 +146,9 @@ impl LinkedBuildData {
     }
 
     /// Fetches target bytecode from linked contracts.
-    pub fn get_target_contract(&self) -> Result<ContractData> {
+    pub fn get_target_contract(&self) -> Result<&ContractData> {
         self.known_contracts
             .get(&self.build_data.target)
-            .cloned()
             .ok_or_eyre("target not found in linked artifacts")
     }
 }

crates/script/src/execute.rs

@@ -73,13 +73,13 @@ impl LinkedState {
             args,
             script_config,
             script_wallets,
-            build_data,
             execution_data: ExecutionData {
                 func,
                 calldata,
                 bytecode: bytecode.clone(),
-                abi: target_contract.abi,
+                abi: target_contract.abi.clone(),
             },
+            build_data,
         })
     }
 }

crates/script/src/lib.rs

@@ -44,7 +44,7 @@ use foundry_evm::{
 };
 use foundry_wallets::MultiWalletOpts;
 use serde::{Deserialize, Serialize};
-use std::{collections::HashMap, sync::Arc};
+use std::collections::HashMap;
 use yansi::Paint;
 
 mod broadcast;
@@ -591,7 +591,7 @@ impl ScriptConfig {
                         CheatsConfig::new(
                             &self.config,
                             self.evm_opts.clone(),
-                            Some(Arc::new(known_contracts)),
+                            Some(known_contracts),
                             Some(script_wallets),
                             Some(target.version),
                         )