chore: TestSetup cleanup

crates/evm/coverage/src/lib.rs

@@ -174,6 +174,14 @@ impl CoverageReport {
 pub struct HitMaps(pub HashMap<B256, HitMap>);
 
 impl HitMaps {
+    pub fn merge_opt(a: &mut Option<Self>, b: Option<Self>) {
+        match (a, b) {
+            (_, None) => {}
+            (a @ None, Some(b)) => *a = Some(b),
+            (Some(a), Some(b)) => a.merge(b),
+        }
+    }
+
     pub fn merge(&mut self, other: Self) {
         for (code_hash, hit_map) in other.0 {
             if let Some(HitMap { hits: extra_hits, .. }) = self.insert(code_hash, hit_map) {

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

@@ -127,11 +127,7 @@ impl FuzzedExecutor {
                         data.logs.extend(case.logs);
                     }
 
-                    // Collect and merge coverage if `forge snapshot` context.
-                    match &mut data.coverage {
-                        Some(prev) => prev.merge(case.coverage.unwrap()),
-                        opt => *opt = case.coverage,
-                    }
+                    HitMaps::merge_opt(&mut data.coverage, case.coverage);
 
                     data.deprecated_cheatcodes = case.deprecated_cheatcodes;
 

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

@@ -204,10 +204,7 @@ impl InvariantTest {
 
     /// Merge current collected coverage with the new coverage from last fuzzed call.
     pub fn merge_coverage(&self, new_coverage: Option<HitMaps>) {
-        match &mut self.execution_data.borrow_mut().coverage {
-            Some(prev) => prev.merge(new_coverage.unwrap()),
-            opt => *opt = new_coverage,
-        }
+        HitMaps::merge_opt(&mut self.execution_data.borrow_mut().coverage, new_coverage);
     }
 
     /// Update metrics for a fuzzed selector, extracted from tx details.

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

@@ -48,16 +48,10 @@ pub fn replay_run(
             tx.call_details.calldata.clone(),
             U256::ZERO,
         )?;
+
         logs.extend(call_result.logs);
         traces.push((TraceKind::Execution, call_result.traces.clone().unwrap()));
-
-        if let Some(new_coverage) = call_result.coverage {
-            if let Some(old_coverage) = coverage {
-                *coverage = Some(std::mem::take(old_coverage).merged(new_coverage));
-            } else {
-                *coverage = Some(new_coverage);
-            }
-        }
+        HitMaps::merge_opt(coverage, call_result.coverage);
 
         // Identify newly generated contracts, if they exist.
         ided_contracts

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

@@ -713,6 +713,12 @@ impl std::ops::DerefMut for DeployResult {
     }
 }
 
+impl From<DeployResult> for RawCallResult {
+    fn from(d: DeployResult) -> Self {
+        d.raw
+    }
+}
+
 /// The result of a raw call.
 #[derive(Debug)]
 pub struct RawCallResult {
@@ -780,6 +786,23 @@ impl Default for RawCallResult {
 }
 
 impl RawCallResult {
+    /// Unpacks an EVM result.
+    pub fn from_evm_result(r: Result<Self, EvmError>) -> eyre::Result<(Self, Option<String>)> {
+        match r {
+            Ok(r) => Ok((r, None)),
+            Err(EvmError::Execution(e)) => Ok((e.raw, Some(e.reason))),
+            Err(e) => Err(e.into()),
+        }
+    }
+
+    /// Unpacks an execution result.
+    pub fn from_execution_result(r: Result<Self, ExecutionErr>) -> (Self, Option<String>) {
+        match r {
+            Ok(r) => (r, None),
+            Err(e) => (e.raw, Some(e.reason)),
+        }
+    }
+
     /// Converts the result of the call into an `EvmError`.
     pub fn into_evm_error(self, rd: Option<&RevertDecoder>) -> EvmError {
         if let Some(reason) = SkipReason::decode(&self.result) {

crates/forge/src/result.rs

@@ -13,7 +13,7 @@ use foundry_common::{evm::Breakpoints, get_contract_name, get_file_name, shell};
 use foundry_evm::{
     coverage::HitMaps,
     decode::SkipReason,
-    executors::{invariant::InvariantMetrics, EvmError, RawCallResult},
+    executors::{invariant::InvariantMetrics, RawCallResult},
     fuzz::{CounterExample, FuzzCase, FuzzFixtures, FuzzTestResult},
     traces::{CallTraceArena, CallTraceDecoder, TraceKind, Traces},
 };
@@ -469,7 +469,7 @@ impl TestResult {
     /// Creates a new test result starting from test setup results.
     pub fn new(setup: TestSetup) -> Self {
         Self {
-            labeled_addresses: setup.labeled_addresses,
+            labeled_addresses: setup.labels,
             logs: setup.logs,
             traces: setup.traces,
             coverage: setup.coverage,
@@ -490,7 +490,7 @@ impl TestResult {
             logs: setup.logs,
             traces: setup.traces,
             coverage: setup.coverage,
-            labeled_addresses: setup.labeled_addresses,
+            labeled_addresses: setup.labels,
             ..Default::default()
         }
     }
@@ -651,21 +651,17 @@ impl TestResult {
         format!("{self} {name} {}", self.kind.report())
     }
 
-    /// Function to merge logs, addresses, traces and coverage from a call result into test result.
-    pub fn merge_call_result(&mut self, call_result: &RawCallResult) {
-        self.logs.extend(call_result.logs.clone());
-        self.labeled_addresses.extend(call_result.labels.clone());
-        self.traces.extend(call_result.traces.clone().map(|traces| (TraceKind::Execution, traces)));
-        self.merge_coverages(call_result.coverage.clone());
+    /// Merges the given raw call result into `self`.
+    pub fn extend(&mut self, call_result: RawCallResult) {
+        self.logs.extend(call_result.logs);
+        self.labeled_addresses.extend(call_result.labels);
+        self.traces.extend(call_result.traces.map(|traces| (TraceKind::Execution, traces)));
+        self.merge_coverages(call_result.coverage);
     }
 
-    /// Function to merge given coverage in current test result coverage.
+    /// Merges the given coverage result into `self`.
     pub fn merge_coverages(&mut self, other_coverage: Option<HitMaps>) {
-        let old_coverage = std::mem::take(&mut self.coverage);
-        self.coverage = match (old_coverage, other_coverage) {
-            (Some(old_coverage), Some(other)) => Some(old_coverage.merged(other)),
-            (a, b) => a.or(b),
-        };
+        HitMaps::merge_opt(&mut self.coverage, other_coverage);
     }
 }
 
@@ -747,77 +743,39 @@ impl TestKind {
     }
 }
 
+/// The result of a test setup.
+///
+/// Includes the deployment of the required libraries and the test contract itself, and the call to
+/// the `setUp()` function.
 #[derive(Clone, Debug, Default)]
 pub struct TestSetup {
-    /// The address at which the test contract was deployed
+    /// The address at which the test contract was deployed.
     pub address: Address,
-    /// The logs emitted during setup
+    /// Defined fuzz test fixtures.
+    pub fuzz_fixtures: FuzzFixtures,
+
+    /// The logs emitted during setup.
     pub logs: Vec<Log>,
-    /// Call traces of the setup
+    /// Addresses labeled during setup.
+    pub labels: AddressHashMap<String>,
+    /// Call traces of the setup.
     pub traces: Traces,
-    /// Addresses labeled during setup
-    pub labeled_addresses: AddressHashMap<String>,
-    /// The reason the setup failed, if it did
-    pub reason: Option<String>,
-    /// Coverage info during setup
+    /// Coverage info during setup.
     pub coverage: Option<HitMaps>,
-    /// Defined fuzz test fixtures
-    pub fuzz_fixtures: FuzzFixtures,
+
+    /// The reason the setup failed, if it did.
+    pub reason: Option<String>,
 }
 
 impl TestSetup {
-    pub fn from_evm_error_with(
-        error: EvmError,
-        mut logs: Vec<Log>,
-        mut traces: Traces,
-        mut labeled_addresses: AddressHashMap<String>,
-    ) -> Self {
-        match error {
-            EvmError::Execution(err) => {
-                // force the tracekind to be setup so a trace is shown.
-                traces.extend(err.raw.traces.map(|traces| (TraceKind::Setup, traces)));
-                logs.extend(err.raw.logs);
-                labeled_addresses.extend(err.raw.labels);
-                Self::failed_with(logs, traces, labeled_addresses, err.reason)
-            }
-            e => Self::failed_with(
-                logs,
-                traces,
-                labeled_addresses,
-                format!("failed to deploy contract: {e}"),
-            ),
-        }
-    }
-
-    pub fn success(
-        address: Address,
-        logs: Vec<Log>,
-        traces: Traces,
-        labeled_addresses: AddressHashMap<String>,
-        coverage: Option<HitMaps>,
-        fuzz_fixtures: FuzzFixtures,
-    ) -> Self {
-        Self { address, logs, traces, labeled_addresses, reason: None, coverage, fuzz_fixtures }
-    }
-
-    pub fn failed_with(
-        logs: Vec<Log>,
-        traces: Traces,
-        labeled_addresses: AddressHashMap<String>,
-        reason: String,
-    ) -> Self {
-        Self {
-            address: Address::ZERO,
-            logs,
-            traces,
-            labeled_addresses,
-            reason: Some(reason),
-            coverage: None,
-            fuzz_fixtures: FuzzFixtures::default(),
-        }
-    }
-
     pub fn failed(reason: String) -> Self {
         Self { reason: Some(reason), ..Default::default() }
     }
+
+    pub fn extend(&mut self, raw: RawCallResult, trace_kind: TraceKind) {
+        self.logs.extend(raw.logs);
+        self.labels.extend(raw.labels);
+        self.traces.extend(raw.traces.map(|traces| (trace_kind, traces)));
+        HitMaps::merge_opt(&mut self.coverage, raw.coverage);
+    }
 }