foundry_evm/executors/invariant/
shrink.rs1use 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#[derive(Debug)]
17struct CallSequenceShrinker {
18 call_sequence_len: usize,
20 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 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 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 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 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 Ok((false, _)) if shrinker.included_calls.count() == 1 => break,
83 Ok((true, _)) => shrinker.included_calls.set(call_idx),
85 _ => {}
86 }
87
88 if let Some(progress) = progress {
89 progress.inc(1);
90 }
91
92 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
103pub 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 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 if call_result.reverted && fail_on_revert && call_result.result.as_ref() != MAGIC_ASSUME {
127 return Ok((false, false));
130 }
131 }
132
133 let (_, mut success) = call_invariant_function(&executor, test_address, calldata)?;
135 if success && call_after_invariant {
138 (_, success) = call_after_invariant_function(&executor, test_address)?;
139 }
140
141 Ok((success, true))
142}