fix(invariants): support vm.assume in invariant tests (#7309)

* fix(invariants): support vm.assume in invariant tests

* fix

* add .sol file

* review fix

Author: klkvr

crates/config/src/invariant.rs

@@ -32,6 +32,9 @@ pub struct InvariantConfig {
     /// Useful for handlers that use cheatcodes as roll or warp
     /// Use it with caution, introduces performance penalty.
     pub preserve_state: bool,
+    /// The maximum number of rejects via `vm.assume` which can be encountered during a single
+    /// invariant run.
+    pub max_assume_rejects: u32,
 }
 
 impl Default for InvariantConfig {
@@ -45,6 +48,7 @@ impl Default for InvariantConfig {
             shrink_sequence: true,
             shrink_run_limit: 2usize.pow(18_u32),
             preserve_state: false,
+            max_assume_rejects: 65536,
         }
     }
 }

crates/evm/evm/src/executors/invariant/error.rs

@@ -53,7 +53,27 @@ pub struct InvariantFuzzTestResult {
 }
 
 #[derive(Clone, Debug)]
-pub struct InvariantFuzzError {
+pub enum InvariantFuzzError {
+    Revert(FailedInvariantCaseData),
+    BrokenInvariant(FailedInvariantCaseData),
+    MaxAssumeRejects(u32),
+}
+
+impl InvariantFuzzError {
+    pub fn revert_reason(&self) -> Option<String> {
+        match self {
+            Self::BrokenInvariant(case_data) | Self::Revert(case_data) => {
+                (!case_data.revert_reason.is_empty()).then(|| case_data.revert_reason.clone())
+            }
+            Self::MaxAssumeRejects(allowed) => Some(format!(
+                "The `vm.assume` cheatcode rejected too many inputs ({allowed} allowed)"
+            )),
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct FailedInvariantCaseData {
     pub logs: Vec<Log>,
     pub traces: Option<CallTraceArena>,
     /// The proptest error occurred as a result of a test case.
@@ -74,7 +94,7 @@ pub struct InvariantFuzzError {
     pub shrink_run_limit: usize,
 }
 
-impl InvariantFuzzError {
+impl FailedInvariantCaseData {
     pub fn new(
         invariant_contract: &InvariantContract<'_>,
         error_func: Option<&Function>,
@@ -93,7 +113,7 @@ impl InvariantFuzzError {
             .with_abi(invariant_contract.abi)
             .decode(call_result.result.as_ref(), Some(call_result.exit_reason));
 
-        InvariantFuzzError {
+        Self {
             logs: call_result.logs,
             traces: call_result.traces,
             test_error: proptest::test_runner::TestError::Fail(

crates/evm/evm/src/executors/invariant/funcs.rs

@@ -1,4 +1,4 @@
-use super::{InvariantFailures, InvariantFuzzError};
+use super::{error::FailedInvariantCaseData, InvariantFailures, InvariantFuzzError};
 use crate::executors::{Executor, RawCallResult};
 use alloy_dyn_abi::JsonAbiExt;
 use alloy_json_abi::Function;
@@ -50,15 +50,16 @@ pub fn assert_invariants(
     if is_err {
         // We only care about invariants which we haven't broken yet.
         if invariant_failures.error.is_none() {
-            invariant_failures.error = Some(InvariantFuzzError::new(
+            let case_data = FailedInvariantCaseData::new(
                 invariant_contract,
                 Some(func),
                 calldata,
                 call_result,
                 &inner_sequence,
                 shrink_sequence,
                 shrink_run_limit,
-            ));
+            );
+            invariant_failures.error = Some(InvariantFuzzError::BrokenInvariant(case_data));
             return None
         }
     }

crates/evm/evm/src/executors/invariant/mod.rs

@@ -9,7 +9,7 @@ use eyre::{eyre, ContextCompat, Result};
 use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact};
 use foundry_config::{FuzzDictionaryConfig, InvariantConfig};
 use foundry_evm_core::{
-    constants::{CALLER, CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS},
+    constants::{CALLER, CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS, MAGIC_ASSUME},
     utils::{get_function, StateChangeset},
 };
 use foundry_evm_fuzz::{
@@ -38,11 +38,13 @@ use foundry_evm_fuzz::strategies::CalldataFuzzDictionary;
 mod funcs;
 pub use funcs::{assert_invariants, replay_run};
 
+use self::error::FailedInvariantCaseData;
+
 /// Alias for (Dictionary for fuzzing, initial contracts to fuzz and an InvariantStrategy).
 type InvariantPreparation = (
     EvmFuzzState,
     FuzzRunIdentifiedContracts,
-    BoxedStrategy<Vec<BasicTxDetails>>,
+    BoxedStrategy<BasicTxDetails>,
     CalldataFuzzDictionary,
 );
 
@@ -143,7 +145,9 @@ impl<'a> InvariantExecutor<'a> {
         // during the run. We need another proptest runner to query for random
         // values.
         let branch_runner = RefCell::new(self.runner.clone());
-        let _ = self.runner.run(&strat, |mut inputs| {
+        let _ = self.runner.run(&strat, |first_input| {
+            let mut inputs = vec![first_input];
+
             // We stop the run immediately if we have reverted, and `fail_on_revert` is set.
             if self.config.fail_on_revert && failures.borrow().reverts > 0 {
                 return Err(TestCaseError::fail("Revert occurred."))
@@ -158,7 +162,10 @@ impl<'a> InvariantExecutor<'a> {
             // Created contracts during a run.
             let mut created_contracts = vec![];
 
-            for current_run in 0..self.config.depth {
+            let mut current_run = 0;
+            let mut assume_rejects_counter = 0;
+
+            while current_run < self.config.depth {
                 let (sender, (address, calldata)) = inputs.last().expect("no input generated");
 
                 // Executes the call from the randomly generated sequence.
@@ -172,65 +179,77 @@ impl<'a> InvariantExecutor<'a> {
                         .expect("could not make raw evm call")
                 };
 
-                // Collect data for fuzzing from the state changeset.
-                let mut state_changeset =
-                    call_result.state_changeset.to_owned().expect("no changesets");
-
-                collect_data(
-                    &mut state_changeset,
-                    sender,
-                    &call_result,
-                    fuzz_state.clone(),
-                    &self.config.dictionary,
-                );
+                if call_result.result.as_ref() == MAGIC_ASSUME {
+                    inputs.pop();
+                    assume_rejects_counter += 1;
+                    if assume_rejects_counter > self.config.max_assume_rejects {
+                        failures.borrow_mut().error = Some(InvariantFuzzError::MaxAssumeRejects(
+                            self.config.max_assume_rejects,
+                        ));
+                        return Err(TestCaseError::fail("Max number of vm.assume rejects reached."))
+                    }
+                } else {
+                    // Collect data for fuzzing from the state changeset.
+                    let mut state_changeset =
+                        call_result.state_changeset.to_owned().expect("no changesets");
+
+                    collect_data(
+                        &mut state_changeset,
+                        sender,
+                        &call_result,
+                        fuzz_state.clone(),
+                        &self.config.dictionary,
+                    );
 
-                if let Err(error) = collect_created_contracts(
-                    &state_changeset,
-                    self.project_contracts,
-                    self.setup_contracts,
-                    &self.artifact_filters,
-                    targeted_contracts.clone(),
-                    &mut created_contracts,
-                ) {
-                    warn!(target: "forge::test", "{error}");
-                }
+                    if let Err(error) = collect_created_contracts(
+                        &state_changeset,
+                        self.project_contracts,
+                        self.setup_contracts,
+                        &self.artifact_filters,
+                        targeted_contracts.clone(),
+                        &mut created_contracts,
+                    ) {
+                        warn!(target: "forge::test", "{error}");
+                    }
 
-                // Commit changes to the database.
-                executor.backend.commit(state_changeset.clone());
-
-                fuzz_runs.push(FuzzCase {
-                    calldata: calldata.clone(),
-                    gas: call_result.gas_used,
-                    stipend: call_result.stipend,
-                });
-
-                let RichInvariantResults { success: can_continue, call_result: call_results } =
-                    can_continue(
-                        &invariant_contract,
-                        call_result,
-                        &executor,
-                        &inputs,
-                        &mut failures.borrow_mut(),
-                        &targeted_contracts,
-                        state_changeset,
-                        self.config.fail_on_revert,
-                        self.config.shrink_sequence,
-                        self.config.shrink_run_limit,
-                    );
+                    // Commit changes to the database.
+                    executor.backend.commit(state_changeset.clone());
+
+                    fuzz_runs.push(FuzzCase {
+                        calldata: calldata.clone(),
+                        gas: call_result.gas_used,
+                        stipend: call_result.stipend,
+                    });
+
+                    let RichInvariantResults { success: can_continue, call_result: call_results } =
+                        can_continue(
+                            &invariant_contract,
+                            call_result,
+                            &executor,
+                            &inputs,
+                            &mut failures.borrow_mut(),
+                            &targeted_contracts,
+                            state_changeset,
+                            self.config.fail_on_revert,
+                            self.config.shrink_sequence,
+                            self.config.shrink_run_limit,
+                        );
+
+                    if !can_continue || current_run == self.config.depth - 1 {
+                        *last_run_calldata.borrow_mut() = inputs.clone();
+                    }
 
-                if !can_continue || current_run == self.config.depth - 1 {
-                    *last_run_calldata.borrow_mut() = inputs.clone();
-                }
+                    if !can_continue {
+                        break
+                    }
 
-                if !can_continue {
-                    break
+                    *last_call_results.borrow_mut() = call_results;
+                    current_run += 1;
                 }
 
-                *last_call_results.borrow_mut() = call_results;
-
                 // Generates the next call from the run using the recently updated
                 // dictionary.
-                inputs.extend(
+                inputs.push(
                     strat
                         .new_tree(&mut branch_runner.borrow_mut())
                         .map_err(|_| TestCaseError::Fail("Could not generate case".into()))?
@@ -772,7 +791,7 @@ fn can_continue(
         failures.reverts += 1;
         // If fail on revert is set, we must return immediately.
         if fail_on_revert {
-            let error = InvariantFuzzError::new(
+            let case_data = FailedInvariantCaseData::new(
                 invariant_contract,
                 None,
                 calldata,
@@ -781,8 +800,8 @@ fn can_continue(
                 shrink_sequence,
                 shrink_run_limit,
             );
-
-            failures.revert_reason = Some(error.revert_reason.clone());
+            failures.revert_reason = Some(case_data.revert_reason.clone());
+            let error = InvariantFuzzError::Revert(case_data);
             failures.error = Some(error);
 
             return RichInvariantResults::new(false, None)

crates/evm/fuzz/src/strategies/invariants.rs

@@ -59,11 +59,10 @@ pub fn invariant_strat(
     contracts: FuzzRunIdentifiedContracts,
     dictionary_weight: u32,
     calldata_fuzz_config: CalldataFuzzDictionary,
-) -> impl Strategy<Value = Vec<BasicTxDetails>> {
+) -> impl Strategy<Value = BasicTxDetails> {
     // We only want to seed the first value, since we want to generate the rest as we mutate the
     // state
     generate_call(fuzz_state, senders, contracts, dictionary_weight, calldata_fuzz_config)
-        .prop_map(|x| vec![x])
 }
 
 /// Strategy to generate a transaction where the `sender`, `target` and `calldata` are all generated

crates/forge/src/runner.rs

@@ -25,7 +25,7 @@ use foundry_evm::{
     fuzz::{invariant::InvariantContract, CounterExample},
     traces::{load_contracts, TraceKind},
 };
-use proptest::test_runner::{TestError, TestRunner};
+use proptest::test_runner::TestRunner;
 use rayon::prelude::*;
 use std::{
     collections::{BTreeMap, HashMap},
@@ -513,26 +513,28 @@ impl<'a> ContractRunner<'a> {
         let mut logs = logs.clone();
         let mut traces = traces.clone();
         let success = error.is_none();
-        let reason = error
-            .as_ref()
-            .and_then(|err| (!err.revert_reason.is_empty()).then(|| err.revert_reason.clone()));
+        let reason = error.as_ref().and_then(|err| err.revert_reason());
         let mut coverage = coverage.clone();
         match error {
             // If invariants were broken, replay the error to collect logs and traces
-            Some(error @ InvariantFuzzError { test_error: TestError::Fail(_, _), .. }) => {
-                match error.replay(
-                    self.executor.clone(),
-                    known_contracts,
-                    identified_contracts.clone(),
-                    &mut logs,
-                    &mut traces,
-                ) {
-                    Ok(c) => counterexample = c,
-                    Err(err) => {
-                        error!(%err, "Failed to replay invariant error");
-                    }
-                };
-            }
+            Some(error) => match error {
+                InvariantFuzzError::BrokenInvariant(case_data) |
+                InvariantFuzzError::Revert(case_data) => {
+                    match case_data.replay(
+                        self.executor.clone(),
+                        known_contracts,
+                        identified_contracts.clone(),
+                        &mut logs,
+                        &mut traces,
+                    ) {
+                        Ok(c) => counterexample = c,
+                        Err(err) => {
+                            error!(%err, "Failed to replay invariant error");
+                        }
+                    };
+                }
+                InvariantFuzzError::MaxAssumeRejects(_) => {}
+            },
 
             // If invariants ran successfully, replay the last run to collect logs and
             // traces.