foundry_cheatcodes/test/
assume.rs

1use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Error, Result};
2use alloy_primitives::Address;
3use foundry_evm_core::constants::MAGIC_ASSUME;
4use spec::Vm::{
5    PotentialRevert, assumeCall, assumeNoRevert_0Call, assumeNoRevert_1Call, assumeNoRevert_2Call,
6};
7use std::fmt::Debug;
8
9#[derive(Clone, Debug)]
10pub struct AssumeNoRevert {
11    /// The call depth at which the cheatcode was added.
12    pub depth: usize,
13    /// Acceptable revert parameters for the next call, to be thrown out if they are encountered;
14    /// reverts with parameters not specified here will count as normal reverts and not rejects
15    /// towards the counter.
16    pub reasons: Vec<AcceptableRevertParameters>,
17    /// Address that reverted the call.
18    pub reverted_by: Option<Address>,
19}
20
21/// Parameters for a single anticipated revert, to be thrown out if encountered.
22#[derive(Clone, Debug)]
23pub struct AcceptableRevertParameters {
24    /// The expected revert data returned by the revert
25    pub reason: Vec<u8>,
26    /// If true then only the first 4 bytes of expected data returned by the revert are checked.
27    pub partial_match: bool,
28    /// Contract expected to revert next call.
29    pub reverter: Option<Address>,
30}
31
32impl AcceptableRevertParameters {
33    fn from(potential_revert: &PotentialRevert) -> Self {
34        Self {
35            reason: potential_revert.revertData.to_vec(),
36            partial_match: potential_revert.partialMatch,
37            reverter: if potential_revert.reverter == Address::ZERO {
38                None
39            } else {
40                Some(potential_revert.reverter)
41            },
42        }
43    }
44}
45
46impl Cheatcode for assumeCall {
47    fn apply(&self, _state: &mut Cheatcodes) -> Result {
48        let Self { condition } = self;
49        if *condition { Ok(Default::default()) } else { Err(Error::from(MAGIC_ASSUME)) }
50    }
51}
52
53impl Cheatcode for assumeNoRevert_0Call {
54    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
55        assume_no_revert(ccx.state, ccx.ecx.journaled_state.depth, vec![])
56    }
57}
58
59impl Cheatcode for assumeNoRevert_1Call {
60    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
61        let Self { potentialRevert } = self;
62        assume_no_revert(
63            ccx.state,
64            ccx.ecx.journaled_state.depth,
65            vec![AcceptableRevertParameters::from(potentialRevert)],
66        )
67    }
68}
69
70impl Cheatcode for assumeNoRevert_2Call {
71    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
72        let Self { potentialReverts } = self;
73        assume_no_revert(
74            ccx.state,
75            ccx.ecx.journaled_state.depth,
76            potentialReverts.iter().map(AcceptableRevertParameters::from).collect(),
77        )
78    }
79}
80
81fn assume_no_revert(
82    state: &mut Cheatcodes,
83    depth: usize,
84    parameters: Vec<AcceptableRevertParameters>,
85) -> Result {
86    ensure!(
87        state.assume_no_revert.is_none(),
88        "you must make another external call prior to calling assumeNoRevert again"
89    );
90
91    state.assume_no_revert = Some(AssumeNoRevert { depth, reasons: parameters, reverted_by: None });
92
93    Ok(Default::default())
94}