fix(forge): add coverage to test setUp (#6123)

* fix(forge): add coverage to test setUp

* Update crates/forge/src/runner.rs

---------

Co-authored-by: evalir <hi@enriqueortiz.dev>

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

@@ -5,6 +5,7 @@ use alloy_json_abi::Function;
 use ethers::types::Log;
 use foundry_common::{ContractsByAddress, ContractsByArtifact};
 use foundry_evm_core::constants::CALLER;
+use foundry_evm_coverage::HitMaps;
 use foundry_evm_fuzz::invariant::{BasicTxDetails, InvariantContract};
 use foundry_evm_traces::{load_contracts, TraceKind, Traces};
 use revm::primitives::U256;
@@ -72,6 +73,7 @@ pub fn replay_run(
     mut ided_contracts: ContractsByAddress,
     logs: &mut Vec<Log>,
     traces: &mut Traces,
+    coverage: &mut Option<HitMaps>,
     func: Function,
     inputs: Vec<BasicTxDetails>,
 ) {
@@ -89,6 +91,20 @@ pub fn replay_run(
         logs.extend(call_result.logs);
         traces.push((TraceKind::Execution, call_result.traces.clone().unwrap()));
 
+        let old_coverage = std::mem::take(coverage);
+        match (old_coverage, call_result.coverage) {
+            (Some(old_coverage), Some(call_coverage)) => {
+                *coverage = Some(old_coverage.merge(call_coverage));
+            }
+            (None, Some(call_coverage)) => {
+                *coverage = Some(call_coverage);
+            }
+            (Some(old_coverage), None) => {
+                *coverage = Some(old_coverage);
+            }
+            (None, None) => {}
+        }
+
         // Identify newly generated contracts, if they exist.
         ided_contracts.extend(load_contracts(
             vec![(TraceKind::Execution, call_result.traces.clone().unwrap())],

crates/evm/evm/src/executors/mod.rs

@@ -366,6 +366,7 @@ impl Executor {
             debug,
             script_wallets,
             env,
+            coverage,
             ..
         } = result;
 
@@ -421,7 +422,7 @@ impl Executor {
 
         trace!(address=?address, "deployed contract");
 
-        Ok(DeployResult { address, gas_used, gas_refunded, logs, traces, debug, env })
+        Ok(DeployResult { address, gas_used, gas_refunded, logs, traces, debug, env, coverage })
     }
 
     /// Deploys a contract and commits the new state to the underlying database.
@@ -597,6 +598,8 @@ pub struct DeployResult {
     pub debug: Option<DebugArena>,
     /// The `revm::Env` after deployment
     pub env: Env,
+    /// The coverage info collected during the deployment
+    pub coverage: Option<HitMaps>,
 }
 
 /// The result of a call.

crates/forge/src/result.rs

@@ -230,6 +230,8 @@ pub struct TestSetup {
     pub labeled_addresses: BTreeMap<Address, String>,
     /// The reason the setup failed, if it did
     pub reason: Option<String>,
+    /// Coverage info during setup
+    pub coverage: Option<HitMaps>,
 }
 
 impl TestSetup {
@@ -261,8 +263,9 @@ impl TestSetup {
         logs: Vec<Log>,
         traces: Traces,
         labeled_addresses: BTreeMap<Address, String>,
+        coverage: Option<HitMaps>,
     ) -> Self {
-        Self { address, logs, traces, labeled_addresses, reason: None }
+        Self { address, logs, traces, labeled_addresses, reason: None, coverage }
     }
 
     pub fn failed_with(
@@ -271,7 +274,14 @@ impl TestSetup {
         labeled_addresses: BTreeMap<Address, String>,
         reason: String,
     ) -> Self {
-        Self { address: Address::ZERO, logs, traces, labeled_addresses, reason: Some(reason) }
+        Self {
+            address: Address::ZERO,
+            logs,
+            traces,
+            labeled_addresses,
+            reason: Some(reason),
+            coverage: None,
+        }
     }
 
     pub fn failed(reason: String) -> Self {

crates/forge/src/runner.rs

@@ -12,6 +12,7 @@ use foundry_common::{
 use foundry_config::{FuzzConfig, InvariantConfig};
 use foundry_evm::{
     constants::CALLER,
+    coverage::HitMaps,
     decode::decode_console_logs,
     executors::{
         fuzz::{CaseOutcome, CounterExampleOutcome, FuzzOutcome, FuzzedExecutor},
@@ -136,28 +137,30 @@ impl<'a> ContractRunner<'a> {
         // Optionally call the `setUp` function
         let setup = if setup {
             trace!("setting up");
-            let (setup_logs, setup_traces, labeled_addresses, reason) =
-                match self.executor.setup(None, address) {
-                    Ok(CallResult { traces, labels, logs, .. }) => {
-                        trace!(contract = ?address, "successfully setUp test");
-                        (logs, traces, labels, None)
-                    }
-                    Err(EvmError::Execution(err)) => {
-                        let ExecutionErr { traces, labels, logs, reason, .. } = *err;
-                        error!(reason = ?reason, contract = ?address, "setUp failed");
-                        (logs, traces, labels, Some(format!("Setup failed: {reason}")))
-                    }
-                    Err(err) => {
-                        error!(reason=?err, contract= ?address, "setUp failed");
-                        (Vec::new(), None, BTreeMap::new(), Some(format!("Setup failed: {err}")))
-                    }
-                };
+            let (setup_logs, setup_traces, labeled_addresses, reason, coverage) = match self
+                .executor
+                .setup(None, address)
+            {
+                Ok(CallResult { traces, labels, logs, coverage, .. }) => {
+                    trace!(contract = ?address, "successfully setUp test");
+                    (logs, traces, labels, None, coverage)
+                }
+                Err(EvmError::Execution(err)) => {
+                    let ExecutionErr { traces, labels, logs, reason, .. } = *err;
+                    error!(reason = ?reason, contract = ?address, "setUp failed");
+                    (logs, traces, labels, Some(format!("Setup failed: {reason}")), None)
+                }
+                Err(err) => {
+                    error!(reason=?err, contract= ?address, "setUp failed");
+                    (Vec::new(), None, BTreeMap::new(), Some(format!("Setup failed: {err}")), None)
+                }
+            };
             traces.extend(setup_traces.map(|traces| (TraceKind::Setup, traces)));
             logs.extend(setup_logs);
 
-            TestSetup { address, logs, traces, labeled_addresses, reason }
+            TestSetup { address, logs, traces, labeled_addresses, reason, coverage }
         } else {
-            TestSetup::success(address, logs, traces, Default::default())
+            TestSetup::success(address, logs, traces, Default::default(), None)
         };
 
         Ok(setup)
@@ -225,7 +228,7 @@ impl<'a> ContractRunner<'a> {
                         logs: setup.logs,
                         kind: TestKind::Standard(0),
                         traces: setup.traces,
-                        coverage: None,
+                        coverage: setup.coverage,
                         labeled_addresses: setup.labeled_addresses,
                         ..Default::default()
                     },
@@ -297,7 +300,9 @@ impl<'a> ContractRunner<'a> {
     /// similar to `eth_call`.
     #[instrument(name = "test", skip_all, fields(name = %func.signature(), %should_fail))]
     pub fn run_test(&self, func: &Function, should_fail: bool, setup: TestSetup) -> TestResult {
-        let TestSetup { address, mut logs, mut traces, mut labeled_addresses, .. } = setup;
+        let TestSetup {
+            address, mut logs, mut traces, mut labeled_addresses, mut coverage, ..
+        } = setup;
 
         // Run unit test
         let mut executor = self.executor.clone();
@@ -318,7 +323,7 @@ impl<'a> ContractRunner<'a> {
                     stipend,
                     logs: execution_logs,
                     traces: execution_trace,
-                    coverage,
+                    coverage: execution_coverage,
                     labels: new_labels,
                     state_changeset,
                     debug,
@@ -329,6 +334,8 @@ impl<'a> ContractRunner<'a> {
                     labeled_addresses.extend(new_labels);
                     logs.extend(execution_logs);
                     debug_arena = debug;
+                    coverage = merge_coverages(coverage, execution_coverage);
+
                     (reverted, None, gas, stipend, coverage, state_changeset, breakpoints)
                 }
                 Err(EvmError::Execution(err)) => {
@@ -415,7 +422,7 @@ impl<'a> ContractRunner<'a> {
         trace!(target: "forge::test::fuzz", "executing invariant test for {:?}", func.name);
         let empty = ContractsByArtifact::default();
         let project_contracts = known_contracts.unwrap_or(&empty);
-        let TestSetup { address, logs, traces, labeled_addresses, .. } = setup;
+        let TestSetup { address, logs, traces, labeled_addresses, coverage, .. } = setup;
 
         // First, run the test normally to see if it needs to be skipped.
         if let Err(EvmError::SkipError) = self.executor.clone().execute_test::<_, _>(
@@ -433,6 +440,7 @@ impl<'a> ContractRunner<'a> {
                 traces,
                 labeled_addresses,
                 kind: TestKind::Invariant { runs: 1, calls: 1, reverts: 1 },
+                coverage,
                 ..Default::default()
             }
         };
@@ -472,6 +480,7 @@ impl<'a> ContractRunner<'a> {
         let reason = error
             .as_ref()
             .and_then(|err| (!err.revert_reason.is_empty()).then(|| err.revert_reason.clone()));
+        let mut coverage = coverage.clone();
         match error {
             // If invariants were broken, replay the error to collect logs and traces
             Some(error @ InvariantFuzzError { test_error: TestError::Fail(_, _), .. }) => {
@@ -499,6 +508,7 @@ impl<'a> ContractRunner<'a> {
                     identified_contracts.clone(),
                     &mut logs,
                     &mut traces,
+                    &mut coverage,
                     func.clone(),
                     last_run_inputs.clone(),
                 );
@@ -521,7 +531,7 @@ impl<'a> ContractRunner<'a> {
             decoded_logs: decode_console_logs(&logs),
             logs,
             kind,
-            coverage: None, // TODO ?
+            coverage,
             traces,
             labeled_addresses: labeled_addresses.clone(),
             ..Default::default() // TODO collect debug traces on the last run or error
@@ -537,7 +547,9 @@ impl<'a> ContractRunner<'a> {
         setup: TestSetup,
         fuzz_config: FuzzConfig,
     ) -> TestResult {
-        let TestSetup { address, mut logs, mut traces, mut labeled_addresses, .. } = setup;
+        let TestSetup {
+            address, mut logs, mut traces, mut labeled_addresses, mut coverage, ..
+        } = setup;
 
         // Run fuzz test
         let start = Instant::now();
@@ -561,6 +573,7 @@ impl<'a> ContractRunner<'a> {
                 kind: TestKind::Standard(0),
                 debug,
                 breakpoints,
+                coverage,
                 ..Default::default()
             }
         }
@@ -614,6 +627,7 @@ impl<'a> ContractRunner<'a> {
         logs.append(&mut result.logs);
         labeled_addresses.append(&mut result.labeled_addresses);
         traces.extend(result.traces.map(|traces| (TraceKind::Execution, traces)));
+        coverage = merge_coverages(coverage, result.coverage);
 
         // Record test execution time
         debug!(
@@ -632,10 +646,21 @@ impl<'a> ContractRunner<'a> {
             logs,
             kind,
             traces,
-            coverage: result.coverage,
+            coverage,
             labeled_addresses,
             debug,
             breakpoints,
         }
     }
 }
+
+/// Utility function to merge coverage options
+fn merge_coverages(mut coverage: Option<HitMaps>, other: Option<HitMaps>) -> Option<HitMaps> {
+    let old_coverage = std::mem::take(&mut coverage);
+    match (old_coverage, other) {
+        (Some(old_coverage), Some(other)) => Some(old_coverage.merge(other)),
+        (None, Some(other)) => Some(other),
+        (Some(old_coverage), None) => Some(old_coverage),
+        (None, None) => None,
+    }
+}

crates/forge/tests/cli/coverage.rs

@@ -1,4 +1,5 @@
 use foundry_test_utils::{forgetest, TestCommand, TestProject};
+use regex::Regex;
 
 forgetest!(basic_coverage, |_prj: TestProject, mut cmd: TestCommand| {
     cmd.args(["coverage"]);
@@ -14,3 +15,74 @@ forgetest!(report_file_coverage, |prj: TestProject, mut cmd: TestCommand| {
     ]);
     cmd.assert_success();
 });
+
+forgetest!(test_setup_coverage, |prj: TestProject, mut cmd: TestCommand| {
+    prj.insert_ds_test();
+    prj.inner()
+        .add_source(
+            "AContract.sol",
+            r#"
+// SPDX-license-identifier: MIT
+pragma solidity ^0.8.0;
+
+contract AContract {
+    int public i;
+
+    function init() public {
+        i = 0;
+    }
+
+    function foo() public {
+        i = 1;
+    }
+}
+    "#,
+        )
+        .unwrap();
+
+    prj.inner()
+        .add_source(
+            "AContractTest.sol",
+            r#"
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.10;
+import "./test.sol";
+import {AContract} from "./AContract.sol";
+
+contract AContractTest is DSTest {
+    AContract a;
+
+    function setUp() public {
+        a = new AContract();
+        a.init();
+    }
+
+    function testFoo() public {
+        a.foo();
+    }
+}
+    "#,
+        )
+        .unwrap();
+
+    let lcov_info = prj.root().join("lcov.info");
+    cmd.arg("coverage").args([
+        "--report".to_string(),
+        "lcov".to_string(),
+        "--report-file".to_string(),
+        lcov_info.to_str().unwrap().to_string(),
+    ]);
+    cmd.assert_success();
+    assert!(lcov_info.exists());
+
+    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)));
+});