fix(invariant): do not commit state if assume returns (#9062)

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

@@ -317,10 +317,11 @@ impl<'a> InvariantExecutor<'a> {
                     TestCaseError::fail("No input generated to call fuzzed target.")
                 })?;
 
-                // Execute call from the randomly generated sequence and commit state changes.
-                let call_result = current_run
+                // Execute call from the randomly generated sequence without committing state.
+                // State is committed only if call is not a magic assume.
+                let mut call_result = current_run
                     .executor
-                    .transact_raw(
+                    .call_raw(
                         tx.sender,
                         tx.call_details.target,
                         tx.call_details.calldata.clone(),
@@ -343,9 +344,11 @@ impl<'a> InvariantExecutor<'a> {
                         return Err(TestCaseError::fail("Max number of vm.assume rejects reached."))
                     }
                 } else {
+                    // Commit executed call result.
+                    current_run.executor.commit(&mut call_result);
+
                     // Collect data for fuzzing from the state changeset.
                     let mut state_changeset = call_result.state_changeset.clone();
-
                     if !call_result.reverted {
                         collect_data(
                             &invariant_test,
@@ -369,13 +372,13 @@ impl<'a> InvariantExecutor<'a> {
                     {
                         warn!(target: "forge::test", "{error}");
                     }
-
                     current_run.fuzz_runs.push(FuzzCase {
                         calldata: tx.call_details.calldata.clone(),
                         gas: call_result.gas_used,
                         stipend: call_result.stipend,
                     });
 
+                    // Determine if test can continue or should exit.
                     let result = can_continue(
                         &invariant_contract,
                         &invariant_test,
@@ -385,11 +388,9 @@ impl<'a> InvariantExecutor<'a> {
                         &state_changeset,
                     )
                     .map_err(|e| TestCaseError::fail(e.to_string()))?;
-
                     if !result.can_continue || current_run.depth == self.config.depth - 1 {
                         invariant_test.set_last_run_inputs(&current_run.inputs);
                     }
-
                     // If test cannot continue then stop current run and exit test suite.
                     if !result.can_continue {
                         return Err(TestCaseError::fail("Test cannot continue."))

crates/forge/tests/it/invariant.rs

@@ -769,3 +769,57 @@ contract AssumeTest is Test {
 ...
 "#]]);
 });
+
+// Test too many inputs rejected for `assumePrecompile`/`assumeForgeAddress`.
+// <https://github.com/foundry-rs/foundry/issues/9054>
+forgetest_init!(should_revert_with_assume_code, |prj, cmd| {
+    let config = Config {
+        invariant: {
+            InvariantConfig { fail_on_revert: true, max_assume_rejects: 10, ..Default::default() }
+        },
+        ..Default::default()
+    };
+    prj.write_config(config);
+
+    // Add initial test that breaks invariant.
+    prj.add_test(
+        "AssumeTest.t.sol",
+        r#"
+import {Test} from "forge-std/Test.sol";
+
+contract BalanceTestHandler is Test {
+    address public ref = address(1412323);
+    address alice;
+
+    constructor(address _alice) {
+        alice = _alice;
+    }
+
+    function increment(uint256 amount_, address addr) public {
+        assumeNotPrecompile(addr);
+        assumeNotForgeAddress(addr);
+        assertEq(alice.balance, 100_000 ether);
+    }
+}
+
+contract BalanceAssumeTest is Test {
+    function setUp() public {
+        address alice = makeAddr("alice");
+        vm.deal(alice, 100_000 ether);
+        targetSender(alice);
+        BalanceTestHandler handler = new BalanceTestHandler(alice);
+        targetContract(address(handler));
+    }
+
+    function invariant_balance() public {}
+}
+     "#,
+    )
+    .unwrap();
+
+    cmd.args(["test", "--mt", "invariant_balance"]).assert_failure().stdout_eq(str![[r#"
+...
+[FAIL: `vm.assume` rejected too many inputs (10 allowed)] invariant_balance() (runs: 0, calls: 0, reverts: 0)
+...
+"#]]);
+});