feat(forge): support for console.log on Invariant handlers (foundry-rs#5488)

* feat: record logs and traces of the last call for all invariants

* feat: fill logs and traces with last call returned logs/traces

* chore: simplify return type

* chore: use struct instead of tuple

* chore: show all invariant logs across all depths when failing

* chore: insert all logs, even if the function reverts

* chore: heavily simplify types

* chore: clippy

* chore: fmt

Author: Evalir

evm/src/fuzz/invariant/error.rs

@@ -145,6 +145,9 @@ impl InvariantFuzzError {
                     .call_raw(CALLER, self.addr, func.0.clone(), U256::zero())
                     .expect("bad call to evm");
 
+                traces.push((TraceKind::Execution, error_call_result.traces.clone().unwrap()));
+
+                logs.extend(error_call_result.logs);
                 if error_call_result.reverted {
                     break
                 }

evm/src/fuzz/invariant/executor.rs

@@ -41,6 +41,20 @@ use std::{cell::RefCell, collections::BTreeMap, sync::Arc};
 type InvariantPreparation =
     (EvmFuzzState, FuzzRunIdentifiedContracts, BoxedStrategy<Vec<BasicTxDetails>>);
 
+/// Enriched results of an invariant run check.
+///
+/// Contains the success condition and call results of the last run
+struct RichInvariantResults {
+    success: bool,
+    call_results: Option<BTreeMap<String, RawCallResult>>,
+}
+
+impl RichInvariantResults {
+    fn new(success: bool, call_results: Option<BTreeMap<String, RawCallResult>>) -> Self {
+        Self { success, call_results }
+    }
+}
+
 /// Wrapper around any [`Executor`] implementor which provides fuzzing support using [`proptest`](https://docs.rs/proptest/1.0.0/proptest/).
 ///
 /// After instantiation, calling `fuzz` will proceed to hammer the deployed smart contracts with
@@ -107,6 +121,7 @@ impl<'a> InvariantExecutor<'a> {
             )
             .ok(),
         );
+        let last_run_calldata: RefCell<Vec<BasicTxDetails>> = RefCell::new(vec![]);
         // Make sure invariants are sound even before starting to fuzz
         if last_call_results.borrow().is_none() {
             fuzz_cases.borrow_mut().push(FuzzedCases::new(vec![]));
@@ -142,7 +157,7 @@ impl<'a> InvariantExecutor<'a> {
                 // Created contracts during a run.
                 let mut created_contracts = vec![];
 
-                'fuzz_run: for _ in 0..self.config.depth {
+                'fuzz_run: for current_run in 0..self.config.depth {
                     let (sender, (address, calldata)) =
                         inputs.last().expect("to have the next randomly generated input.");
 
@@ -183,7 +198,7 @@ impl<'a> InvariantExecutor<'a> {
                         stipend: call_result.stipend,
                     });
 
-                    let (can_continue, call_results) = can_continue(
+                    let RichInvariantResults { success: can_continue, call_results } = can_continue(
                         &invariant_contract,
                         call_result,
                         &executor,
@@ -195,6 +210,10 @@ impl<'a> InvariantExecutor<'a> {
                         self.config.shrink_sequence,
                     );
 
+                    if !can_continue || current_run == self.config.depth - 1 {
+                        *last_run_calldata.borrow_mut() = inputs.clone();
+                    }
+
                     if !can_continue {
                         break 'fuzz_run
                     }
@@ -233,7 +252,7 @@ impl<'a> InvariantExecutor<'a> {
             invariants,
             cases: fuzz_cases.into_inner(),
             reverts,
-            last_call_results: last_call_results.take(),
+            last_run_inputs: last_run_calldata.take(),
         })
     }
 
@@ -556,7 +575,8 @@ fn collect_data(
 }
 
 /// Verifies that the invariant run execution can continue.
-/// Returns the mapping of (Invariant Function Name -> Call Result) if invariants were asserted.
+/// Returns the mapping of (Invariant Function Name -> Call Result, Logs, Traces) if invariants were
+/// asserted.
 #[allow(clippy::too_many_arguments)]
 fn can_continue(
     invariant_contract: &InvariantContract,
@@ -568,7 +588,7 @@ fn can_continue(
     state_changeset: StateChangeset,
     fail_on_revert: bool,
     shrink_sequence: bool,
-) -> (bool, Option<BTreeMap<String, RawCallResult>>) {
+) -> RichInvariantResults {
     let mut call_results = None;
 
     // Detect handler assertion failures first.
@@ -583,7 +603,7 @@ fn can_continue(
             assert_invariants(invariant_contract, executor, calldata, failures, shrink_sequence)
                 .ok();
         if call_results.is_none() {
-            return (false, None)
+            return RichInvariantResults::new(false, None)
         }
     } else {
         failures.reverts += 1;
@@ -604,13 +624,16 @@ fn can_continue(
 
             // Hacky to provide the full error to the user.
             for invariant in invariant_contract.invariant_functions.iter() {
-                failures.failed_invariants.insert(invariant.name.clone(), Some(error.clone()));
+                failures.failed_invariants.insert(
+                    invariant.name.clone(),
+                    (Some(error.clone()), invariant.to_owned().clone()),
+                );
             }
 
-            return (false, None)
+            return RichInvariantResults::new(false, None)
         }
     }
-    (true, call_results)
+    RichInvariantResults::new(true, call_results)
 }
 
 #[derive(Clone)]
@@ -623,21 +646,24 @@ pub struct InvariantFailures {
     /// How many different invariants have been broken.
     pub broken_invariants_count: usize,
     /// Maps a broken invariant to its specific error.
-    pub failed_invariants: BTreeMap<String, Option<InvariantFuzzError>>,
+    pub failed_invariants: BTreeMap<String, (Option<InvariantFuzzError>, Function)>,
 }
 
 impl InvariantFailures {
     fn new(invariants: &[&Function]) -> Self {
         InvariantFailures {
             reverts: 0,
             broken_invariants_count: 0,
-            failed_invariants: invariants.iter().map(|f| (f.name.to_string(), None)).collect(),
+            failed_invariants: invariants
+                .iter()
+                .map(|f| (f.name.to_string(), (None, f.to_owned().clone())))
+                .collect(),
             revert_reason: None,
         }
     }
 
     /// Moves `reverts` and `failed_invariants` out of the struct.
-    fn into_inner(self) -> (usize, BTreeMap<String, Option<InvariantFuzzError>>) {
+    fn into_inner(self) -> (usize, BTreeMap<String, (Option<InvariantFuzzError>, Function)>) {
         (self.reverts, self.failed_invariants)
     }
 }

evm/src/fuzz/invariant/mod.rs

@@ -1,11 +1,16 @@
 //! Fuzzing support abstracted over the [`Evm`](crate::Evm) used
-use crate::{fuzz::*, CALLER};
+use crate::{
+    fuzz::*,
+    trace::{load_contracts, TraceKind, Traces},
+    CALLER,
+};
 mod error;
 pub use error::InvariantFuzzError;
 mod filters;
 pub use filters::{ArtifactFilters, SenderFilters};
 mod call_override;
 pub use call_override::{set_up_inner_replay, RandomCallGenerator};
+use foundry_common::ContractsByArtifact;
 mod executor;
 use crate::executor::Executor;
 use ethers::{
@@ -87,17 +92,20 @@ pub fn assert_invariants(
                 .expect("to have been initialized.");
 
             // We only care about invariants which we haven't broken yet.
-            if invariant_error.is_none() {
+            if invariant_error.0.is_none() {
                 invariant_failures.failed_invariants.insert(
                     broken_invariant.name.clone(),
-                    Some(InvariantFuzzError::new(
-                        invariant_contract,
-                        Some(broken_invariant),
-                        calldata,
-                        call_result,
-                        &inner_sequence,
-                        shrink_sequence,
-                    )),
+                    (
+                        Some(InvariantFuzzError::new(
+                            invariant_contract,
+                            Some(broken_invariant),
+                            calldata,
+                            call_result,
+                            &inner_sequence,
+                            shrink_sequence,
+                        )),
+                        broken_invariant.clone().to_owned(),
+                    ),
                 );
                 found_case = true;
             } else {
@@ -114,7 +122,7 @@ pub fn assert_invariants(
         invariant_failures.broken_invariants_count = invariant_failures
             .failed_invariants
             .iter()
-            .filter(|(_function, error)| error.is_some())
+            .filter(|(_function, error)| error.0.is_some())
             .count();
 
         eyre::bail!(
@@ -126,14 +134,64 @@ pub fn assert_invariants(
     Ok(call_results)
 }
 
+/// Replays the provided invariant run for collecting the logs and traces from all depths.
+#[allow(clippy::too_many_arguments)]
+pub fn replay_run(
+    invariant_contract: &InvariantContract,
+    mut executor: Executor,
+    known_contracts: Option<&ContractsByArtifact>,
+    mut ided_contracts: ContractsByAddress,
+    logs: &mut Vec<Log>,
+    traces: &mut Traces,
+    func: Function,
+    inputs: Vec<BasicTxDetails>,
+) {
+    // We want traces for a failed case.
+    executor.set_tracing(true);
+
+    // set_up_inner_replay(&mut executor, &inputs);
+
+    // Replay each call from the sequence until we break the invariant.
+    for (sender, (addr, bytes)) in inputs.iter() {
+        let call_result = executor
+            .call_raw_committing(*sender, *addr, bytes.0.clone(), U256::zero())
+            .expect("bad call to evm");
+
+        logs.extend(call_result.logs);
+        traces.push((TraceKind::Execution, call_result.traces.clone().unwrap()));
+
+        // Identify newly generated contracts, if they exist.
+        ided_contracts.extend(load_contracts(
+            vec![(TraceKind::Execution, call_result.traces.clone().unwrap())],
+            known_contracts,
+        ));
+
+        // Checks the invariant.
+        let error_call_result = executor
+            .call_raw(
+                CALLER,
+                invariant_contract.address,
+                func.encode_input(&[]).expect("invariant should have no inputs").into(),
+                U256::zero(),
+            )
+            .expect("bad call to evm");
+
+        traces.push((TraceKind::Execution, error_call_result.traces.clone().unwrap()));
+
+        logs.extend(error_call_result.logs);
+    }
+}
+
 /// The outcome of an invariant fuzz test
 #[derive(Debug)]
 pub struct InvariantFuzzTestResult {
-    pub invariants: BTreeMap<String, Option<InvariantFuzzError>>,
+    pub invariants: BTreeMap<String, (Option<InvariantFuzzError>, Function)>,
     /// Every successful fuzz test case
     pub cases: Vec<FuzzedCases>,
     /// Number of reverted fuzz calls
     pub reverts: usize,
 
-    pub last_call_results: Option<BTreeMap<String, RawCallResult>>,
+    /// The entire inputs of the last run of the invariant campaign, used for
+    /// replaying the run for collecting traces.
+    pub last_run_inputs: Vec<BasicTxDetails>,
 }

forge/src/runner.rs

@@ -17,7 +17,8 @@ use foundry_evm::{
     executor::{CallResult, EvmError, ExecutionErr, Executor},
     fuzz::{
         invariant::{
-            InvariantContract, InvariantExecutor, InvariantFuzzError, InvariantFuzzTestResult,
+            replay_run, InvariantContract, InvariantExecutor, InvariantFuzzError,
+            InvariantFuzzTestResult,
         },
         FuzzedExecutor,
     },
@@ -449,15 +450,16 @@ impl<'a> ContractRunner<'a> {
         let invariant_contract =
             InvariantContract { address, invariant_functions: functions, abi: self.contract };
 
-        let Ok(InvariantFuzzTestResult { invariants, cases, reverts, mut last_call_results }) =
-            evm.invariant_fuzz(invariant_contract)
+        let Ok(InvariantFuzzTestResult { invariants, cases, reverts, last_run_inputs }) =
+            evm.invariant_fuzz(invariant_contract.clone())
         else {
             return vec![]
         };
 
         invariants
-            .into_iter()
-            .map(|(func_name, test_error)| {
+            .into_values()
+            .map(|test_error| {
+                let (test_error, invariant) = test_error;
                 let mut counterexample = None;
                 let mut logs = logs.clone();
                 let mut traces = traces.clone();
@@ -489,18 +491,19 @@ impl<'a> ContractRunner<'a> {
                             traces.push((TraceKind::Execution, error_traces));
                         }
                     }
-                    // If invariants ran successfully, collect last call logs and traces
                     _ => {
-                        if let Some(last_call_result) = last_call_results
-                            .as_mut()
-                            .and_then(|call_results| call_results.remove(&func_name))
-                        {
-                            logs.extend(last_call_result.logs);
-
-                            if let Some(last_call_traces) = last_call_result.traces {
-                                traces.push((TraceKind::Execution, last_call_traces));
-                            }
-                        }
+                        // If invariants ran successfully, replay the last run to collect logs and
+                        // traces.
+                        replay_run(
+                            &invariant_contract.clone(),
+                            self.executor.clone(),
+                            known_contracts,
+                            identified_contracts.clone(),
+                            &mut logs,
+                            &mut traces,
+                            invariant,
+                            last_run_inputs.clone(),
+                        );
                     }
                 }