feat(invariant): shrink failure when replayed (#12351)

Author: grandizzy

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

@@ -30,7 +30,6 @@ use parking_lot::RwLock;
 use proptest::{strategy::Strategy, test_runner::TestRunner};
 use result::{assert_after_invariant, assert_invariants, can_continue};
 use revm::state::Account;
-use shrink::shrink_sequence;
 use std::{
     collections::{HashMap as Map, btree_map::Entry},
     sync::Arc,
@@ -323,6 +322,10 @@ impl<'a> InvariantExecutor<'a> {
         }
     }
 
+    pub fn config(self) -> InvariantConfig {
+        self.config
+    }
+
     /// Fuzzes any deployed contract and checks any broken invariant at `invariant_address`.
     pub fn invariant_fuzz(
         &mut self,

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

@@ -1,18 +1,15 @@
-use super::{
-    call_after_invariant_function, call_invariant_function, error::FailedInvariantCaseData,
-    shrink_sequence,
-};
-use crate::executors::{EarlyExit, Executor};
+use super::{call_after_invariant_function, call_invariant_function};
+use crate::executors::{EarlyExit, Executor, invariant::shrink::shrink_sequence};
 use alloy_dyn_abi::JsonAbiExt;
 use alloy_primitives::{Log, U256, map::HashMap};
 use eyre::Result;
 use foundry_common::{ContractsByAddress, ContractsByArtifact};
+use foundry_config::InvariantConfig;
 use foundry_evm_coverage::HitMaps;
 use foundry_evm_fuzz::{BaseCounterExample, BasicTxDetails, invariant::InvariantContract};
 use foundry_evm_traces::{TraceKind, TraceMode, Traces, load_contracts};
 use indicatif::ProgressBar;
 use parking_lot::RwLock;
-use proptest::test_runner::TestError;
 use std::sync::Arc;
 
 /// Replays a call sequence for collecting logs and traces.
@@ -98,50 +95,41 @@ pub fn replay_run(
 /// Replays the error case, shrinks the failing sequence and collects all necessary traces.
 #[expect(clippy::too_many_arguments)]
 pub fn replay_error(
-    failed_case: &FailedInvariantCaseData,
-    invariant_contract: &InvariantContract<'_>,
+    config: InvariantConfig,
     mut executor: Executor,
+    calls: &[BasicTxDetails],
+    inner_sequence: Option<Vec<Option<BasicTxDetails>>>,
+    invariant_contract: &InvariantContract<'_>,
     known_contracts: &ContractsByArtifact,
     ided_contracts: ContractsByAddress,
     logs: &mut Vec<Log>,
     traces: &mut Traces,
     line_coverage: &mut Option<HitMaps>,
     deprecated_cheatcodes: &mut HashMap<&'static str, Option<&'static str>>,
     progress: Option<&ProgressBar>,
-    show_solidity: bool,
     early_exit: &EarlyExit,
 ) -> Result<Vec<BaseCounterExample>> {
-    match failed_case.test_error {
-        // Don't use at the moment.
-        TestError::Abort(_) => Ok(vec![]),
-        TestError::Fail(_, ref calls) => {
-            // Shrink sequence of failed calls.
-            let calls = shrink_sequence(
-                failed_case,
-                calls,
-                &executor,
-                invariant_contract.call_after_invariant,
-                progress,
-                early_exit,
-            )?;
+    // Shrink sequence of failed calls.
+    let calls =
+        shrink_sequence(&config, invariant_contract, calls, &executor, progress, early_exit)?;
 
-            set_up_inner_replay(&mut executor, &failed_case.inner_sequence);
-
-            // Replay calls to get the counterexample and to collect logs, traces and coverage.
-            replay_run(
-                invariant_contract,
-                executor,
-                known_contracts,
-                ided_contracts,
-                logs,
-                traces,
-                line_coverage,
-                deprecated_cheatcodes,
-                &calls,
-                show_solidity,
-            )
-        }
+    if let Some(sequence) = inner_sequence {
+        set_up_inner_replay(&mut executor, &sequence);
     }
+
+    // Replay calls to get the counterexample and to collect logs, traces and coverage.
+    replay_run(
+        invariant_contract,
+        executor,
+        known_contracts,
+        ided_contracts,
+        logs,
+        traces,
+        line_coverage,
+        deprecated_cheatcodes,
+        &calls,
+        config.show_solidity,
+    )
 }
 
 /// Sets up the calls generated by the internal fuzzer, if they exist.

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

@@ -1,12 +1,11 @@
 use crate::executors::{
     EarlyExit, Executor,
-    invariant::{
-        call_after_invariant_function, call_invariant_function, error::FailedInvariantCaseData,
-    },
+    invariant::{call_after_invariant_function, call_invariant_function},
 };
 use alloy_primitives::{Address, Bytes, U256};
+use foundry_config::InvariantConfig;
 use foundry_evm_core::constants::MAGIC_ASSUME;
-use foundry_evm_fuzz::BasicTxDetails;
+use foundry_evm_fuzz::{BasicTxDetails, invariant::InvariantContract};
 use indicatif::ProgressBar;
 use proptest::bits::{BitSetLike, VarBitSet};
 
@@ -33,39 +32,36 @@ impl CallSequenceShrinker {
     }
 }
 
-/// Shrinks the failure case to its smallest sequence of calls.
-///
-/// The shrunk call sequence always respect the order failure is reproduced as it is tested
-/// top-down.
 pub(crate) fn shrink_sequence(
-    failed_case: &FailedInvariantCaseData,
+    config: &InvariantConfig,
+    invariant_contract: &InvariantContract<'_>,
     calls: &[BasicTxDetails],
     executor: &Executor,
-    call_after_invariant: bool,
     progress: Option<&ProgressBar>,
     early_exit: &EarlyExit,
 ) -> eyre::Result<Vec<BasicTxDetails>> {
     trace!(target: "forge::test", "Shrinking sequence of {} calls.", calls.len());
 
     // Reset run count and display shrinking message.
     if let Some(progress) = progress {
-        progress.set_length(failed_case.shrink_run_limit as usize as u64);
+        progress.set_length(config.shrink_run_limit as u64);
         progress.reset();
         progress.set_message(" Shrink");
     }
 
+    let target_address = invariant_contract.address;
+    let calldata: Bytes = invariant_contract.invariant_function.selector().to_vec().into();
     // Special case test: the invariant is *unsatisfiable* - it took 0 calls to
     // break the invariant -- consider emitting a warning.
-    let (_, success) =
-        call_invariant_function(executor, failed_case.addr, failed_case.calldata.clone())?;
+    let (_, success) = call_invariant_function(executor, target_address, calldata.clone())?;
     if !success {
         return Ok(vec![]);
     }
 
     let mut call_idx = 0;
 
     let mut shrinker = CallSequenceShrinker::new(calls.len());
-    for _ in 0..failed_case.shrink_run_limit {
+    for _ in 0..config.shrink_run_limit {
         if early_exit.should_stop() {
             break;
         }
@@ -77,10 +73,10 @@ pub(crate) fn shrink_sequence(
             executor.clone(),
             calls,
             shrinker.current().collect(),
-            failed_case.addr,
-            failed_case.calldata.clone(),
-            failed_case.fail_on_revert,
-            call_after_invariant,
+            target_address,
+            calldata.clone(),
+            config.fail_on_revert,
+            invariant_contract.call_after_invariant,
         ) {
             // If candidate sequence still fails, shrink until shortest possible.
             Ok((false, _)) if shrinker.included_calls.count() == 1 => break,