feat(cheatcodes): add vm.assumeNoRevert for fuzz tests (#8780)

Author: grandizzy

crates/cheatcodes/assets/cheatcodes.json

@@ -678,6 +678,10 @@ interface Vm {
     #[cheatcode(group = Testing, safety = Safe)]
     function assume(bool condition) external pure;
 
+    /// Discard this run's fuzz inputs and generate new ones if next call reverted.
+    #[cheatcode(group = Testing, safety = Safe)]
+    function assumeNoRevert() external pure;
+
     /// Writes a breakpoint to jump to in the debugger.
     #[cheatcode(group = Testing, safety = Safe)]
     function breakpoint(string calldata char) external;

crates/cheatcodes/spec/src/vm.rs

@@ -9,9 +9,12 @@ use crate::{
     },
     inspector::utils::CommonCreateInput,
     script::{Broadcast, ScriptWallets},
-    test::expect::{
-        self, ExpectedCallData, ExpectedCallTracker, ExpectedCallType, ExpectedEmit,
-        ExpectedRevert, ExpectedRevertKind,
+    test::{
+        assume::AssumeNoRevert,
+        expect::{
+            self, ExpectedCallData, ExpectedCallTracker, ExpectedCallType, ExpectedEmit,
+            ExpectedRevert, ExpectedRevertKind,
+        },
     },
     utils::IgnoredTraces,
     CheatsConfig, CheatsCtxt, DynCheatcode, Error, Result, Vm,
@@ -25,7 +28,7 @@ use foundry_config::Config;
 use foundry_evm_core::{
     abi::Vm::stopExpectSafeMemoryCall,
     backend::{DatabaseExt, RevertDiagnostic},
-    constants::{CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS},
+    constants::{CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS, MAGIC_ASSUME},
     utils::new_evm_with_existing_context,
     InspectorExt,
 };
@@ -294,6 +297,9 @@ pub struct Cheatcodes {
     /// Expected revert information
     pub expected_revert: Option<ExpectedRevert>,
 
+    /// Assume next call can revert and discard fuzz run if it does.
+    pub assume_no_revert: Option<AssumeNoRevert>,
+
     /// Additional diagnostic for reverts
     pub fork_revert_diagnostic: Option<RevertDiagnostic>,
 
@@ -384,6 +390,7 @@ impl Cheatcodes {
             gas_price: Default::default(),
             prank: Default::default(),
             expected_revert: Default::default(),
+            assume_no_revert: Default::default(),
             fork_revert_diagnostic: Default::default(),
             accesses: Default::default(),
             recorded_account_diffs_stack: Default::default(),
@@ -1106,6 +1113,19 @@ impl<DB: DatabaseExt> Inspector<DB> 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 outcome.result.is_revert() {
+                    outcome.result.output = Error::from(MAGIC_ASSUME).abi_encode().into();
+                    return outcome;
+                }
+                // Call didn't revert, reset `assume_no_revert` state.
+                self.assume_no_revert = None;
+            }
+        }
+
         // Handle expected reverts
         if let Some(expected_revert) = &self.expected_revert {
             if ecx.journaled_state.depth() <= expected_revert.depth {

crates/cheatcodes/src/inspector.rs

@@ -3,25 +3,15 @@
 use chrono::DateTime;
 use std::env;
 
-use crate::{Cheatcode, Cheatcodes, CheatsCtxt, DatabaseExt, Error, Result, Vm::*};
+use crate::{Cheatcode, Cheatcodes, CheatsCtxt, DatabaseExt, Result, Vm::*};
 use alloy_primitives::Address;
 use alloy_sol_types::SolValue;
-use foundry_evm_core::constants::{MAGIC_ASSUME, MAGIC_SKIP};
+use foundry_evm_core::constants::MAGIC_SKIP;
 
 pub(crate) mod assert;
+pub(crate) mod assume;
 pub(crate) mod expect;
 
-impl Cheatcode for assumeCall {
-    fn apply(&self, _state: &mut Cheatcodes) -> Result {
-        let Self { condition } = self;
-        if *condition {
-            Ok(Default::default())
-        } else {
-            Err(Error::from(MAGIC_ASSUME))
-        }
-    }
-}
-
 impl Cheatcode for breakpoint_0Call {
     fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
         let Self { char } = self;

crates/cheatcodes/src/test.rs

@@ -0,0 +1,29 @@
+use crate::{Cheatcode, Cheatcodes, CheatsCtxt, Error, Result};
+use foundry_evm_core::{backend::DatabaseExt, constants::MAGIC_ASSUME};
+use spec::Vm::{assumeCall, assumeNoRevertCall};
+use std::fmt::Debug;
+
+#[derive(Clone, Debug)]
+pub struct AssumeNoRevert {
+    /// The call depth at which the cheatcode was added.
+    pub depth: u64,
+}
+
+impl Cheatcode for assumeCall {
+    fn apply(&self, _state: &mut Cheatcodes) -> Result {
+        let Self { condition } = self;
+        if *condition {
+            Ok(Default::default())
+        } else {
+            Err(Error::from(MAGIC_ASSUME))
+        }
+    }
+}
+
+impl Cheatcode for assumeNoRevertCall {
+    fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        ccx.state.assume_no_revert =
+            Some(AssumeNoRevert { depth: ccx.ecx.journaled_state.depth() });
+        Ok(Default::default())
+    }
+}

crates/cheatcodes/src/test/assume.rs

@@ -1834,3 +1834,82 @@ contract CounterTest is DSTest {
 ...
 "#]]);
 });
+
+forgetest_init!(test_assume_no_revert, |prj, cmd| {
+    prj.wipe_contracts();
+    prj.insert_ds_test();
+    prj.insert_vm();
+    prj.clear();
+
+    prj.add_source(
+        "Counter.t.sol",
+        r#"pragma solidity 0.8.24;
+import {Vm} from "./Vm.sol";
+import {DSTest} from "./test.sol";
+contract CounterWithRevert {
+    error CountError();
+    error CheckError();
+
+    function count(uint256 a) public pure returns (uint256) {
+        if (a > 1000 || a < 10) {
+            revert CountError();
+        }
+        return 99999999;
+    }
+    function check(uint256 a) public pure {
+        if (a == 99999999) {
+            revert CheckError();
+        }
+    }
+    function dummy() public pure {}
+}
+
+contract CounterRevertTest is DSTest {
+    Vm vm = Vm(HEVM_ADDRESS);
+
+    function test_assume_no_revert_pass(uint256 a) public {
+        CounterWithRevert counter = new CounterWithRevert();
+        vm.assumeNoRevert();
+        a = counter.count(a);
+        assertEq(a, 99999999);
+    }
+    function test_assume_no_revert_fail_assert(uint256 a) public {
+        CounterWithRevert counter = new CounterWithRevert();
+        vm.assumeNoRevert();
+        a = counter.count(a);
+        // Test should fail on next assertion.
+        assertEq(a, 1);
+    }
+    function test_assume_no_revert_fail_in_2nd_call(uint256 a) public {
+        CounterWithRevert counter = new CounterWithRevert();
+        vm.assumeNoRevert();
+        a = counter.count(a);
+        // Test should revert here (not in scope of `assumeNoRevert` cheatcode).
+        counter.check(a);
+        assertEq(a, 99999999);
+    }
+    function test_assume_no_revert_fail_in_3rd_call(uint256 a) public {
+        CounterWithRevert counter = new CounterWithRevert();
+        vm.assumeNoRevert();
+        a = counter.count(a);
+        // Test `assumeNoRevert` applied to non reverting call should not be available for next reverting call.
+        vm.assumeNoRevert();
+        counter.dummy();
+        // Test will revert here (not in scope of `assumeNoRevert` cheatcode).
+        counter.check(a);
+        assertEq(a, 99999999);
+    }
+}
+     "#,
+    )
+    .unwrap();
+
+    cmd.args(["test"]).with_no_redact().assert_failure().stdout_eq(str![[r#"
+...
+[FAIL. Reason: assertion failed; counterexample: [..]] test_assume_no_revert_fail_assert(uint256) [..]
+[FAIL. Reason: CheckError(); counterexample: [..]] test_assume_no_revert_fail_in_2nd_call(uint256) [..]
+[FAIL. Reason: CheckError(); counterexample: [..]] test_assume_no_revert_fail_in_3rd_call(uint256) [..]
+[PASS] test_assume_no_revert_pass(uint256) (runs: 256, [..])
+...
+"#]]);
+});