feat(cheatcodes): add ability to ignore (multiple) specific and parti…

…al reverts in fuzz and invariant tests (#9179)

* initial pass

add support for multiple reasons, add tests

appease clippy

fix broken tests; fix some assume behavior

remove comment and bad error-surfacing logic

remove redundant param, rename revert.rs, create sol test file

remove unnecessary tests from both test_cmd and AssumeNoRevert.t.sol

use empty vec instead of option<vec>; remove commented test

remove assumeNoPartialRevert; update assumeNoPartialRevert

Simplify test, use snapbox assertion

Redact number of runs

implement assume_no_revert change

* rebase and refactor

* fix tests for overloaded; original failing

* remove erroneous return type

* appease clippy

* allow combining expectRevert with assumeNoRevert

* Apply suggestions from code review

nit

* remove magic string const

* fix error string

* improve invariant selectors weight test

* nit

---------

Co-authored-by: zerosnacks <95942363+zerosnacks@users.noreply.github.com>
Co-authored-by: grandizzy <grandizzy.the.egg@gmail.com>

crates/cheatcodes/spec/src/lib.rs

@@ -88,6 +88,7 @@ impl Cheatcodes<'static> {
                 Vm::DebugStep::STRUCT.clone(),
                 Vm::BroadcastTxSummary::STRUCT.clone(),
                 Vm::SignedDelegation::STRUCT.clone(),
+                Vm::PotentialRevert::STRUCT.clone(),
             ]),
             enums: Cow::Owned(vec![
                 Vm::CallerMode::ENUM.clone(),

crates/cheatcodes/spec/src/vm.rs

@@ -324,6 +324,18 @@ interface Vm {
         address implementation;
     }
 
+    /// Represents a "potential" revert reason from a single subsequent call when using `vm.assumeNoReverts`.
+    /// Reverts that match will result in a FOUNDRY::ASSUME rejection, whereas unmatched reverts will be surfaced
+    /// as normal.
+    struct PotentialRevert {
+        /// The allowed origin of the revert opcode; address(0) allows reverts from any address
+        address reverter;
+        /// When true, only matches on the beginning of the revert data, otherwise, matches on entire revert data
+        bool partialMatch;
+        /// The data to use to match encountered reverts
+        bytes revertData;
+    }
+
     // ======== EVM ========
 
     /// Gets the address for a given private key.
@@ -894,6 +906,14 @@ interface Vm {
     #[cheatcode(group = Testing, safety = Safe)]
     function assumeNoRevert() external pure;
 
+    /// Discard this run's fuzz inputs and generate new ones if next call reverts with the potential revert parameters.
+    #[cheatcode(group = Testing, safety = Safe)]
+    function assumeNoRevert(PotentialRevert calldata potentialRevert) external pure;
+
+    /// Discard this run's fuzz inputs and generate new ones if next call reverts with the any of the potential revert parameters.
+    #[cheatcode(group = Testing, safety = Safe)]
+    function assumeNoRevert(PotentialRevert[] calldata potentialReverts) 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, ExpectedEmitTracker,
             ExpectedRevert, ExpectedRevertKind,
         },
+        revert_handlers,
     },
     utils::IgnoredTraces,
     CheatsConfig, CheatsCtxt, DynCheatcode, Error, Result,
@@ -755,16 +756,14 @@ where {
                 matches!(expected_revert.kind, ExpectedRevertKind::Default)
             {
                 let mut expected_revert = std::mem::take(&mut self.expected_revert).unwrap();
-                let handler_result = expect::handle_expect_revert(
+                return match revert_handlers::handle_expect_revert(
                     false,
                     true,
-                    &mut expected_revert,
+                    &expected_revert,
                     outcome.result.result,
                     outcome.result.output.clone(),
                     &self.config.available_artifacts,
-                );
-
-                return match handler_result {
+                ) {
                     Ok((address, retdata)) => {
                         expected_revert.actual_count += 1;
                         if expected_revert.actual_count < expected_revert.count {
@@ -1287,16 +1286,45 @@ 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.
+        // Handle assume no revert cheatcode.
+        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;
             }
         }
 
@@ -1330,20 +1358,15 @@ impl Inspector<&mut dyn DatabaseExt> for Cheatcodes {
                 };
 
                 if needs_processing {
-                    // Only `remove` the expected revert from state if `expected_revert.count` ==
-                    // `expected_revert.actual_count`
                     let mut expected_revert = std::mem::take(&mut self.expected_revert).unwrap();
-
-                    let handler_result = expect::handle_expect_revert(
+                    return match revert_handlers::handle_expect_revert(
                         cheatcode_call,
                         false,
-                        &mut expected_revert,
+                        &expected_revert,
                         outcome.result.result,
                         outcome.result.output.clone(),
                         &self.config.available_artifacts,
-                    );
-
-                    return match handler_result {
+                    ) {
                         Err(error) => {
                             trace!(expected=?expected_revert, ?error, status=?outcome.result.result, "Expected revert mismatch");
                             outcome.result.result = InstructionResult::Revert;

crates/cheatcodes/src/test.rs

@@ -9,6 +9,7 @@ use foundry_evm_core::constants::MAGIC_SKIP;
 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,46 @@
 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, PotentialRevert,
+};
 use std::fmt::Debug;
 
 #[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 AcceptableRevertParameters {
+    fn from(potential_revert: &PotentialRevert) -> Self {
+        Self {
+            reason: potential_revert.revertData.to_vec(),
+            partial_match: potential_revert.partialMatch,
+            reverter: if potential_revert.reverter == Address::ZERO {
+                None
+            } else {
+                Some(potential_revert.reverter)
+            },
+        }
+    }
 }
 
 impl Cheatcode for assumeCall {
@@ -20,10 +54,45 @@ 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(), vec![])
     }
 }
+
+impl Cheatcode for assumeNoRevert_1Call {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { potentialRevert } = self;
+        assume_no_revert(
+            ccx.state,
+            ccx.ecx.journaled_state.depth(),
+            vec![AcceptableRevertParameters::from(potentialRevert)],
+        )
+    }
+}
+
+impl Cheatcode for assumeNoRevert_2Call {
+    fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result {
+        let Self { potentialReverts } = self;
+        assume_no_revert(
+            ccx.state,
+            ccx.ecx.journaled_state.depth(),
+            potentialReverts.iter().map(AcceptableRevertParameters::from).collect(),
+        )
+    }
+}
+
+fn assume_no_revert(
+    state: &mut Cheatcodes,
+    depth: u64,
+    parameters: Vec<AcceptableRevertParameters>,
+) -> Result {
+    ensure!(
+        state.assume_no_revert.is_none(),
+        "you must make another external call prior to calling assumeNoRevert again"
+    );
+
+    state.assume_no_revert = Some(AssumeNoRevert { depth, reasons: parameters, reverted_by: None });
+
+    Ok(Default::default())
+}