refactor: deploy all libraries when running tests (foundry-rs#8034)

* refactor: deploy all libraries when running tests

* fix test

* fix tests

* fix traces test

* review fixes

* fix

* review fixes

Author: klkvr

crates/evm/core/src/backend/mod.rs

@@ -536,11 +536,6 @@ impl Backend {
     /// This will also grant cheatcode access to the test account
     pub fn set_test_contract(&mut self, acc: Address) -> &mut Self {
         trace!(?acc, "setting test account");
-        // toggle the previous sender
-        if let Some(current) = self.inner.test_contract_address.take() {
-            self.remove_persistent_account(&current);
-            self.revoke_cheatcode_access(&acc);
-        }
 
         self.add_persistent_account(acc);
         self.allow_cheatcode_access(acc);

crates/forge/bin/cmd/coverage.rs

@@ -246,6 +246,8 @@ impl CoverageArgs {
             .set_coverage(true)
             .build(&root, output, env, evm_opts)?;
 
+        let known_contracts = runner.known_contracts.clone();
+
         let outcome = self
             .test
             .run_tests(runner, config.clone(), verbosity, &self.test.filter(&config))
@@ -257,12 +259,10 @@ impl CoverageArgs {
             for result in suite.test_results.values() {
                 let Some(hit_maps) = result.coverage.as_ref() else { continue };
                 for map in hit_maps.0.values() {
-                    if let Some((id, _)) =
-                        suite.known_contracts.find_by_deployed_code(&map.bytecode)
-                    {
+                    if let Some((id, _)) = known_contracts.find_by_deployed_code(&map.bytecode) {
                         hits.push((id, map, true));
                     } else if let Some((id, _)) =
-                        suite.known_contracts.find_by_creation_code(&map.bytecode)
+                        known_contracts.find_by_creation_code(&map.bytecode)
                     {
                         hits.push((id, map, false));
                     }

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

@@ -312,11 +312,12 @@ impl TestArgs {
             *test_pattern = Some(debug_test_pattern.clone());
         }
 
+        let libraries = runner.libraries.clone();
         let outcome = self.run_tests(runner, config, verbosity, &filter).await?;
 
         if should_debug {
             // Get first non-empty suite result. We will have only one such entry
-            let Some((suite_result, test_result)) = outcome
+            let Some((_, test_result)) = outcome
                 .results
                 .iter()
                 .find(|(_, r)| !r.test_results.is_empty())
@@ -328,7 +329,7 @@ impl TestArgs {
             let sources = ContractSources::from_project_output(
                 output_clone.as_ref().unwrap(),
                 project.root(),
-                &suite_result.libraries,
+                &libraries,
             )?;
 
             // Run the debugger.
@@ -376,6 +377,7 @@ impl TestArgs {
         }
 
         let remote_chain_id = runner.evm_opts.get_remote_chain_id().await;
+        let known_contracts = runner.known_contracts.clone();
 
         // Run tests.
         let (tx, rx) = channel::<(String, SuiteResult)>();
@@ -385,38 +387,40 @@ impl TestArgs {
             move || runner.test(&filter, tx)
         });
 
+        // Set up trace identifiers.
+        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.
+        if !self.gas_report {
+            identifier = identifier.with_etherscan(&config, remote_chain_id)?;
+        }
+
+        // Build the trace decoder.
+        let mut builder = CallTraceDecoderBuilder::new()
+            .with_known_contracts(&known_contracts)
+            .with_verbosity(verbosity);
+        // Signatures are of no value for gas reports.
+        if !self.gas_report {
+            builder = builder.with_signature_identifier(SignaturesIdentifier::new(
+                Config::foundry_cache_dir(),
+                config.offline,
+            )?);
+        }
+        let mut decoder = builder.build();
+
         let mut gas_report = self
             .gas_report
             .then(|| GasReport::new(config.gas_reports.clone(), config.gas_reports_ignore.clone()));
 
         let mut outcome = TestOutcome::empty(self.allow_failure);
 
         let mut any_test_failed = false;
-        for (contract_name, mut suite_result) in rx {
+        for (contract_name, suite_result) in rx {
             let tests = &suite_result.test_results;
 
-            // Set up trace identifiers.
-            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.
-            if !self.gas_report {
-                identifier = identifier.with_etherscan(&config, remote_chain_id)?;
-            }
-
-            // Build the trace decoder.
-            let mut builder = CallTraceDecoderBuilder::new()
-                .with_known_contracts(known_contracts)
-                .with_verbosity(verbosity);
-            // Signatures are of no value for gas reports.
-            if !self.gas_report {
-                builder = builder.with_signature_identifier(SignaturesIdentifier::new(
-                    Config::foundry_cache_dir(),
-                    config.offline,
-                )?);
-            }
-            let mut decoder = builder.build();
+            // Clear the addresses and labels from previous test.
+            decoder.clear_addresses();
 
             // We identify addresses if we're going to print *any* trace or gas report.
             let identify_addresses = verbosity >= 3 || self.gas_report || self.debug.is_some();
@@ -520,18 +524,15 @@ 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);
 
             // Stop processing the remaining suites if any test failed and `fail_fast` is set.
             if self.fail_fast && any_test_failed {
                 break;
             }
         }
+        outcome.last_run_decoder = Some(decoder);
         let duration = timer.elapsed();
 
         trace!(target: "forge::test", len=outcome.results.len(), %any_test_failed, "done with results");

crates/forge/src/multi_runner.rs

@@ -1,11 +1,13 @@
 //! Forge test runner for multiple contracts.
 
-use crate::{result::SuiteResult, ContractRunner, TestFilter, TestOptions};
+use crate::{
+    result::SuiteResult, runner::LIBRARY_DEPLOYER, ContractRunner, TestFilter, TestOptions,
+};
 use alloy_json_abi::{Function, JsonAbi};
 use alloy_primitives::{Address, Bytes, U256};
 use eyre::Result;
 use foundry_common::{get_contract_name, ContractsByArtifact, TestFunctionExt};
-use foundry_compilers::{artifacts::Libraries, Artifact, ArtifactId, ProjectCompileOutput, Solc};
+use foundry_compilers::{artifacts::Libraries, Artifact, ArtifactId, ProjectCompileOutput};
 use foundry_config::Config;
 use foundry_evm::{
     backend::Backend, decode::RevertDecoder, executors::ExecutorBuilder, fork::CreateFork,
@@ -27,8 +29,6 @@ use std::{
 pub struct TestContract {
     pub abi: JsonAbi,
     pub bytecode: Bytes,
-    pub libs_to_deploy: Vec<Bytes>,
-    pub libraries: Libraries,
 }
 
 pub type DeployableContracts = BTreeMap<ArtifactId, TestContract>;
@@ -61,8 +61,12 @@ pub struct MultiContractRunner {
     pub test_options: TestOptions,
     /// Whether to enable call isolation
     pub isolation: bool,
-    /// Output of the project compilation
-    pub output: ProjectCompileOutput,
+    /// Known contracts linked with computed library addresses.
+    pub known_contracts: ContractsByArtifact,
+    /// Libraries to deploy.
+    pub libs_to_deploy: Vec<Bytes>,
+    /// Library addresses used to link contracts.
+    pub libraries: Libraries,
 }
 
 impl MultiContractRunner {
@@ -181,17 +185,10 @@ impl MultiContractRunner {
         let identifier = artifact_id.identifier();
         let mut span_name = identifier.as_str();
 
-        let linker = Linker::new(
-            self.config.project_paths::<Solc>().root,
-            self.output.artifact_ids().collect(),
-        );
-        let linked_contracts = linker.get_linked_artifacts(&contract.libraries).unwrap_or_default();
-        let known_contracts = ContractsByArtifact::new(linked_contracts);
-
         let cheats_config = CheatsConfig::new(
             &self.config,
             self.evm_opts.clone(),
-            Some(known_contracts.clone()),
+            Some(self.known_contracts.clone()),
             None,
             Some(artifact_id.version.clone()),
         );
@@ -220,12 +217,13 @@ impl MultiContractRunner {
             &identifier,
             executor,
             contract,
+            &self.libs_to_deploy,
             self.evm_opts.initial_balance,
             self.sender,
             &self.revert_decoder,
             self.debug,
         );
-        let r = runner.run_tests(filter, &self.test_options, known_contracts, handle);
+        let r = runner.run_tests(filter, &self.test_options, self.known_contracts.clone(), handle);
 
         debug!(duration=?r.duration, "executed all tests in contract");
 
@@ -332,10 +330,19 @@ impl MultiContractRunnerBuilder {
             .filter_map(|(_, contract)| contract.abi.as_ref().map(|abi| abi.borrow()));
         let revert_decoder = RevertDecoder::new().with_abis(abis);
 
+        let LinkOutput { libraries, libs_to_deploy } = linker.link_with_nonce_or_address(
+            Default::default(),
+            LIBRARY_DEPLOYER,
+            0,
+            linker.contracts.keys(),
+        )?;
+
+        let linked_contracts = linker.get_linked_artifacts(&libraries)?;
+
         // Create a mapping of name => (abi, deployment code, Vec<library deployment code>)
         let mut deployable_contracts = DeployableContracts::default();
 
-        for (id, contract) in linker.contracts.iter() {
+        for (id, contract) in linked_contracts.iter() {
             let Some(abi) = contract.abi.as_ref() else {
                 continue;
             };
@@ -344,35 +351,19 @@ impl MultiContractRunnerBuilder {
             if abi.constructor.as_ref().map(|c| c.inputs.is_empty()).unwrap_or(true) &&
                 abi.functions().any(|func| func.name.is_test() || func.name.is_invariant_test())
             {
-                let LinkOutput { libs_to_deploy, libraries } = linker.link_with_nonce_or_address(
-                    Default::default(),
-                    evm_opts.sender,
-                    1,
-                    id,
-                )?;
-
-                let linked_contract = linker.link(id, &libraries)?;
-
-                let Some(bytecode) = linked_contract
-                    .get_bytecode_bytes()
-                    .map(|b| b.into_owned())
-                    .filter(|b| !b.is_empty())
+                let Some(bytecode) =
+                    contract.get_bytecode_bytes().map(|b| b.into_owned()).filter(|b| !b.is_empty())
                 else {
                     continue;
                 };
 
-                deployable_contracts.insert(
-                    id.clone(),
-                    TestContract {
-                        abi: abi.clone().into_owned(),
-                        bytecode,
-                        libs_to_deploy,
-                        libraries,
-                    },
-                );
+                deployable_contracts
+                    .insert(id.clone(), TestContract { abi: abi.clone(), bytecode });
             }
         }
 
+        let known_contracts = ContractsByArtifact::new(linked_contracts);
+
         Ok(MultiContractRunner {
             contracts: deployable_contracts,
             evm_opts,
@@ -386,7 +377,9 @@ impl MultiContractRunnerBuilder {
             debug: self.debug,
             test_options: self.test_options.unwrap_or_default(),
             isolation: self.isolation,
-            output,
+            known_contracts,
+            libs_to_deploy,
+            libraries,
         })
     }
 }

crates/forge/src/result.rs

@@ -2,10 +2,7 @@
 
 use crate::gas_report::GasReport;
 use alloy_primitives::{Address, Log};
-use foundry_common::{
-    evm::Breakpoints, get_contract_name, get_file_name, shell, ContractsByArtifact,
-};
-use foundry_compilers::artifacts::Libraries;
+use foundry_common::{evm::Breakpoints, get_contract_name, get_file_name, shell};
 use foundry_evm::{
     coverage::HitMaps,
     debug::DebugArena,
@@ -196,31 +193,15 @@ pub struct SuiteResult {
     pub test_results: BTreeMap<String, TestResult>,
     /// Generated warnings.
     pub warnings: Vec<String>,
-    /// 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: ContractsByArtifact,
 }
 
 impl SuiteResult {
     pub fn new(
         duration: Duration,
         test_results: BTreeMap<String, TestResult>,
         warnings: Vec<String>,
-        libraries: Libraries,
-        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);
-        }
+        Self { duration, test_results, warnings }
     }
 
     /// Returns an iterator over all individual succeeding tests and their names.