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 pub depth: usize,
13 pub reasons: Vec<AcceptableRevertParameters>,
17 pub reverted_by: Option<Address>,
19}
20
21#[derive(Clone, Debug)]
23pub struct AcceptableRevertParameters {
24 pub reason: Vec<u8>,
26 pub partial_match: bool,
28 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}