foundry_cheatcodes/test/
assume.rs

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

foundry_cheatcodes/test/
assume.rs
