foundry_evm/executors/invariant/
shrink.rs

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