foundry_evm/executors/invariant/
shrink.rs

1use crate::executors::{
2    EarlyExit, Executor,
3    invariant::{
4        call_after_invariant_function, call_invariant_function, error::FailedInvariantCaseData,
5    },
6};
7use alloy_primitives::{Address, Bytes, U256};
8use foundry_evm_core::constants::MAGIC_ASSUME;
9use foundry_evm_fuzz::BasicTxDetails;
10use indicatif::ProgressBar;
11use proptest::bits::{BitSetLike, VarBitSet};
12
13/// Shrinker for a call sequence failure.
14/// Iterates sequence call sequence top down and removes calls one by one.
15/// If the failure is still reproducible with removed call then moves to the next one.
16/// If the failure is not reproducible then restore removed call and moves to next one.
17#[derive(Debug)]
18struct CallSequenceShrinker {
19    /// Length of call sequence to be shrunk.
20    call_sequence_len: usize,
21    /// Call ids contained in current shrunk sequence.
22    included_calls: VarBitSet,
23}
24
25impl CallSequenceShrinker {
26    fn new(call_sequence_len: usize) -> Self {
27        Self { call_sequence_len, included_calls: VarBitSet::saturated(call_sequence_len) }
28    }
29
30    /// Return candidate shrink sequence to be tested, by removing ids from original sequence.
31    fn current(&self) -> impl Iterator<Item = usize> + '_ {
32        (0..self.call_sequence_len).filter(|&call_id| self.included_calls.test(call_id))
33    }
34}
35
36/// Shrinks the failure case to its smallest sequence of calls.
37///
38/// The shrunk call sequence always respect the order failure is reproduced as it is tested
39/// top-down.
40pub(crate) fn shrink_sequence(
41    failed_case: &FailedInvariantCaseData,
42    calls: &[BasicTxDetails],
43    executor: &Executor,
44    call_after_invariant: bool,
45    progress: Option<&ProgressBar>,
46    early_exit: &EarlyExit,
47) -> eyre::Result<Vec<BasicTxDetails>> {
48    trace!(target: "forge::test", "Shrinking sequence of {} calls.", calls.len());
49
50    // Reset run count and display shrinking message.
51    if let Some(progress) = progress {
52        progress.set_length(failed_case.shrink_run_limit as usize as u64);
53        progress.reset();
54        progress.set_message(" Shrink");
55    }
56
57    // Special case test: the invariant is *unsatisfiable* - it took 0 calls to
58    // break the invariant -- consider emitting a warning.
59    let (_, success) =
60        call_invariant_function(executor, failed_case.addr, failed_case.calldata.clone())?;
61    if !success {
62        return Ok(vec![]);
63    }
64
65    let mut call_idx = 0;
66
67    let mut shrinker = CallSequenceShrinker::new(calls.len());
68    for _ in 0..failed_case.shrink_run_limit {
69        if early_exit.should_stop() {
70            break;
71        }
72
73        // Remove call at current index.
74        shrinker.included_calls.clear(call_idx);
75
76        match check_sequence(
77            executor.clone(),
78            calls,
79            shrinker.current().collect(),
80            failed_case.addr,
81            failed_case.calldata.clone(),
82            failed_case.fail_on_revert,
83            call_after_invariant,
84        ) {
85            // If candidate sequence still fails, shrink until shortest possible.
86            Ok((false, _)) if shrinker.included_calls.count() == 1 => break,
87            // Restore last removed call as it caused sequence to pass invariant.
88            Ok((true, _)) => shrinker.included_calls.set(call_idx),
89            _ => {}
90        }
91
92        if let Some(progress) = progress {
93            progress.inc(1);
94        }
95
96        // Restart from first call once we reach the end of sequence.
97        if call_idx + 1 == shrinker.call_sequence_len {
98            call_idx = 0;
99        } else {
100            call_idx += 1;
101        };
102    }
103
104    Ok(shrinker.current().map(|idx| &calls[idx]).cloned().collect())
105}
106
107/// Checks if the given call sequence breaks the invariant.
108///
109/// Used in shrinking phase for checking candidate sequences and in replay failures phase to test
110/// persisted failures.
111/// Returns the result of invariant check (and afterInvariant call if needed) and if sequence was
112/// entirely applied.
113pub fn check_sequence(
114    mut executor: Executor,
115    calls: &[BasicTxDetails],
116    sequence: Vec<usize>,
117    test_address: Address,
118    calldata: Bytes,
119    fail_on_revert: bool,
120    call_after_invariant: bool,
121) -> eyre::Result<(bool, bool)> {
122    // Apply the call sequence.
123    for call_index in sequence {
124        let tx = &calls[call_index];
125        let call_result = executor.transact_raw(
126            tx.sender,
127            tx.call_details.target,
128            tx.call_details.calldata.clone(),
129            U256::ZERO,
130        )?;
131        // Ignore calls reverted with `MAGIC_ASSUME`. This is needed to handle failed scenarios that
132        // are replayed with a modified version of test driver (that use new `vm.assume`
133        // cheatcodes).
134        if call_result.reverted && fail_on_revert && call_result.result.as_ref() != MAGIC_ASSUME {
135            // Candidate sequence fails test.
136            // We don't have to apply remaining calls to check sequence.
137            return Ok((false, false));
138        }
139    }
140
141    // Check the invariant for call sequence.
142    let (_, mut success) = call_invariant_function(&executor, test_address, calldata)?;
143    // Check after invariant result if invariant is success and `afterInvariant` function is
144    // declared.
145    if success && call_after_invariant {
146        (_, success) = call_after_invariant_function(&executor, test_address)?;
147    }
148
149    Ok((success, true))
150}