feat(cheatcodes): add ability to ignore (multiple) specific and partial reverts in fuzz and invariant tests

crates/cheatcodes/spec/src/vm.rs

@@ -825,6 +825,22 @@ interface Vm {
     #[cheatcode(group = Testing, safety = Safe)]
     function assumeNoRevert() external pure;
 
+    /// Discard this run's fuzz inputs and generate new ones if next call reverted with data starting with the given revert data. Call more than once to add multiple reasons.
+    #[cheatcode(group = Testing, safety = Safe)]
+    function assumeNoRevert(bytes4 revertData) external pure;
+
+    /// Discard this run's fuzz inputs and generate new ones if next call reverted with the given revert data. Call more than once to add multiple reasons.
+    #[cheatcode(group = Testing, safety = Safe)]
+    function assumeNoRevert(bytes calldata revertData) external pure;
+
+    /// Discard this run's fuzz inputs and generate new ones if next call reverted with data starting with the given revert data. Call more than once to add multiple reasons.
+    #[cheatcode(group = Testing, safety = Safe)]
+    function assumeNoRevert(bytes4 revertData, address reverter) external pure;
+
+    /// Discard this run's fuzz inputs and generate new ones if next call reverted with the given revert data. Call more than once to add multiple reasons.
+    #[cheatcode(group = Testing, safety = Safe)]
+    function assumeNoRevert(bytes calldata revertData, address reverter) external pure;
+
     /// Writes a breakpoint to jump to in the debugger.
     #[cheatcode(group = Testing, safety = Safe)]
     function breakpoint(string calldata char) external pure;

crates/cheatcodes/src/inspector.rs

@@ -15,6 +15,7 @@ use crate::{
             self, ExpectedCallData, ExpectedCallTracker, ExpectedCallType, ExpectedEmit,
             ExpectedRevert, ExpectedRevertKind,
         },
+        revert_handlers,
     },
     utils::IgnoredTraces,
     CheatsConfig, CheatsCtxt, DynCheatcode, Error, Result,
@@ -746,7 +747,7 @@ where {
                 matches!(expected_revert.kind, ExpectedRevertKind::Default)
             {
                 let expected_revert = std::mem::take(&mut self.expected_revert).unwrap();
-                return match expect::handle_expect_revert(
+                return match revert_handlers::handle_expect_revert(
                     false,
                     true,
                     &expected_revert,
@@ -1252,15 +1253,44 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes {
         }
 
         // Handle assume not revert cheatcode.
-        if let Some(assume_no_revert) = &self.assume_no_revert {
-            if ecx.journaled_state.depth() == assume_no_revert.depth && !cheatcode_call {
-                // Discard run if we're at the same depth as cheatcode and call reverted.
+        if let Some(assume_no_revert) = &mut self.assume_no_revert {
+            // Record current reverter address before processing the expect revert if call reverted,
+            // expect revert is set with expected reverter address and no actual reverter set yet.
+            if outcome.result.is_revert() && assume_no_revert.reverted_by.is_none() {
+                assume_no_revert.reverted_by = Some(call.target_address);
+            }
+            // allow multiple cheatcode calls at the same depth
+            if ecx.journaled_state.depth() <= assume_no_revert.depth && !cheatcode_call {
+                // Discard run if we're at the same depth as cheatcode, call reverted, and no
+                // specific reason was supplied
                 if outcome.result.is_revert() {
-                    outcome.result.output = Error::from(MAGIC_ASSUME).abi_encode().into();
+                    let assume_no_revert = std::mem::take(&mut self.assume_no_revert).unwrap();
+                    return match revert_handlers::handle_assume_no_revert(
+                        &assume_no_revert,
+                        outcome.result.result,
+                        &outcome.result.output,
+                        &self.config.available_artifacts,
+                    ) {
+                        // if result is Ok, it was an anticipated revert; return an "assume" error
+                        // to reject this run
+                        Ok(_) => {
+                            outcome.result.output = Error::from(MAGIC_ASSUME).abi_encode().into();
+                            outcome
+                        }
+                        // if result is Error, it was an unanticipated revert; should revert
+                        // normally
+                        Err(error) => {
+                            trace!(expected=?assume_no_revert, ?error, status=?outcome.result.result, "Expected revert mismatch");
+                            outcome.result.result = InstructionResult::Revert;
+                            outcome.result.output = error.abi_encode().into();
+                            outcome
+                        }
+                    }
+                } else {
+                    // Call didn't revert, reset `assume_no_revert` state.
+                    self.assume_no_revert = None;
                     return outcome;
                 }
-                // Call didn't revert, reset `assume_no_revert` state.
-                self.assume_no_revert = None;
             }
         }
 
@@ -1274,7 +1304,6 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes {
             {
                 expected_revert.reverted_by = Some(call.target_address);
             }
-
             if ecx.journaled_state.depth() <= expected_revert.depth {
                 let needs_processing = match expected_revert.kind {
                     ExpectedRevertKind::Default => !cheatcode_call,
@@ -1287,7 +1316,7 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes {
 
                 if needs_processing {
                     let expected_revert = std::mem::take(&mut self.expected_revert).unwrap();
-                    return match expect::handle_expect_revert(
+                    return match revert_handlers::handle_expect_revert(
                         cheatcode_call,
                         false,
                         &expected_revert,

crates/cheatcodes/src/test.rs

@@ -10,6 +10,7 @@ use std::env;
 pub(crate) mod assert;
 pub(crate) mod assume;
 pub(crate) mod expect;
+pub(crate) mod revert_handlers;
 
 impl Cheatcode for breakpoint_0Call {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {

crates/cheatcodes/src/test/assume.rs

@@ -1,12 +1,37 @@
 use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Error, Result};
+use alloy_primitives::Address;
 use foundry_evm_core::constants::MAGIC_ASSUME;
-use spec::Vm::{assumeCall, assumeNoRevertCall};
+use spec::Vm::{
+    assumeCall, assumeNoRevert_0Call, assumeNoRevert_1Call, assumeNoRevert_2Call,
+    assumeNoRevert_3Call, assumeNoRevert_4Call,
+};
 use std::fmt::Debug;
 
+pub const ASSUME_EXPECT_REJECT_MAGIC: &str = "Cannot combine an assumeNoRevert with expectRevert";
+pub const ASSUME_REJECT_MAGIC: &str =
+    "Cannot combine a generic assumeNoRevert with specific assumeNoRevert reasons";
+
 #[derive(Clone, Debug)]
 pub struct AssumeNoRevert {
     /// The call depth at which the cheatcode was added.
     pub depth: u64,
+    /// Acceptable revert parameters for the next call, to be thrown out if they are encountered;
+    /// reverts with parameters not specified here will count as normal reverts and not rejects
+    /// towards the counter.
+    pub reasons: Vec<AcceptableRevertParameters>,
+    /// Address that reverted the call.
+    pub reverted_by: Option<Address>,
+}
+
+/// Parameters for a single anticipated revert, to be thrown out if encountered.
+#[derive(Clone, Debug)]
+pub struct AcceptableRevertParameters {
+    /// The expected revert data returned by the revert
+    pub reason: Vec<u8>,
+    /// If true then only the first 4 bytes of expected data returned by the revert are checked.
+    pub partial_match: bool,
+    /// Contract expected to revert next call.
+    pub reverter: Option<Address>,
 }
 
 impl Cheatcode for assumeCall {
@@ -20,10 +45,83 @@ impl Cheatcode for assumeCall {
     }
 }
 
-impl Cheatcode for assumeNoRevertCall {
+impl Cheatcode for assumeNoRevert_0Call {
     fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
-        ccx.state.assume_no_revert =
-            Some(AssumeNoRevert { depth: ccx.ecx.journaled_state.depth() });
-        Ok(Default::default())
+        assume_no_revert(ccx.state, ccx.ecx.journaled_state.depth(), None, None)
     }
 }
+
+impl Cheatcode for assumeNoRevert_1Call {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { revertData } = self;
+        assume_no_revert(
+            ccx.state,
+            ccx.ecx.journaled_state.depth(),
+            Some(revertData.to_vec()),
+            None,
+        )
+    }
+}
+impl Cheatcode for assumeNoRevert_2Call {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { revertData } = self;
+        assume_no_revert(
+            ccx.state,
+            ccx.ecx.journaled_state.depth(),
+            Some(revertData.to_vec()),
+            None,
+        )
+    }
+}
+impl Cheatcode for assumeNoRevert_3Call {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { revertData, reverter } = self;
+        assume_no_revert(
+            ccx.state,
+            ccx.ecx.journaled_state.depth(),
+            Some(revertData.to_vec()),
+            Some(*reverter),
+        )
+    }
+}
+impl Cheatcode for assumeNoRevert_4Call {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { revertData, reverter } = self;
+        assume_no_revert(
+            ccx.state,
+            ccx.ecx.journaled_state.depth(),
+            Some(revertData.to_vec()),
+            Some(*reverter),
+        )
+    }
+}
+
+fn assume_no_revert(
+    state: &mut Cheatcodes,
+    depth: u64,
+    reason: Option<Vec<u8>>,
+    reverter: Option<Address>,
+) -> Result {
+    ensure!(state.expected_revert.is_none(), ASSUME_EXPECT_REJECT_MAGIC);
+
+    let params = reason.map(|reason| {
+        let partial_match = reason.len() == 4;
+        AcceptableRevertParameters { reason, partial_match, reverter }
+    });
+
+    match state.assume_no_revert {
+        Some(ref mut assume) => {
+            ensure!(!assume.reasons.is_empty() && params.is_some(), ASSUME_REJECT_MAGIC);
+            assume.reasons.push(params.unwrap());
+        }
+        None => {
+            state.assume_no_revert = Some(AssumeNoRevert {
+                depth,
+                reasons: if let Some(params) = params { vec![params] } else { vec![] },
+                reverted_by: None,
+            });
+        }
+    }
+
+    Ok(Default::default())
+}