fix(forge): preserve invariant replay failure reason (#14136)

* fix(forge): preserve invariant replay failure reason

* test(forge): add invariant replay reason regression tests

* test(forge): add replay test for invariant-level custom error

* fix(forge): decode invariant replay reason with revert decoder

* refactor(evm): reduce check_sequence args for clippy

* test(forge): relax invariant replay snapshots

Author: 0xWeakSheep

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

@@ -55,7 +55,7 @@ mod result;
 pub use result::InvariantFuzzTestResult;
 
 mod shrink;
-pub use shrink::{check_sequence, check_sequence_value};
+pub use shrink::{CheckSequenceOptions, check_sequence, check_sequence_value};
 
 sol! {
     interface IInvariantTest {

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

@@ -1,10 +1,12 @@
 use crate::executors::{
-    EarlyExit, Executor,
+    EarlyExit, EvmError, Executor, RawCallResult,
     invariant::{call_after_invariant_function, call_invariant_function, execute_tx},
 };
 use alloy_primitives::{Address, Bytes, I256, U256};
 use foundry_config::InvariantConfig;
-use foundry_evm_core::{FoundryBlock, constants::MAGIC_ASSUME, evm::FoundryEvmNetwork};
+use foundry_evm_core::{
+    FoundryBlock, constants::MAGIC_ASSUME, decode::RevertDecoder, evm::FoundryEvmNetwork,
+};
 use foundry_evm_fuzz::{BasicTxDetails, invariant::InvariantContract};
 use indicatif::ProgressBar;
 use proptest::bits::{BitSetLike, VarBitSet};
@@ -122,13 +124,16 @@ pub(crate) fn shrink_sequence<FEN: FoundryEvmNetwork>(
             shrinker.current().collect(),
             target_address,
             calldata.clone(),
-            config.fail_on_revert,
-            invariant_contract.call_after_invariant,
+            CheckSequenceOptions {
+                fail_on_revert: config.fail_on_revert,
+                call_after_invariant: invariant_contract.call_after_invariant,
+                rd: None,
+            },
         ) {
             // If candidate sequence still fails, shrink until shortest possible.
-            Ok((false, _)) if shrinker.included_calls.count() == 1 => break,
+            Ok((false, _, _)) if shrinker.included_calls.count() == 1 => break,
             // Restore last removed call as it caused sequence to pass invariant.
-            Ok((true, _)) => shrinker.included_calls.set(call_idx),
+            Ok((true, _, _)) => shrinker.included_calls.set(call_idx),
             _ => {}
         }
 
@@ -154,9 +159,8 @@ pub fn check_sequence<FEN: FoundryEvmNetwork>(
     sequence: Vec<usize>,
     test_address: Address,
     calldata: Bytes,
-    fail_on_revert: bool,
-    call_after_invariant: bool,
-) -> eyre::Result<(bool, bool)> {
+    options: CheckSequenceOptions<'_>,
+) -> eyre::Result<(bool, bool, Option<String>)> {
     // Apply the call sequence.
     for call_index in sequence {
         let tx = &calls[call_index];
@@ -165,22 +169,51 @@ pub fn check_sequence<FEN: FoundryEvmNetwork>(
         // Ignore calls reverted with `MAGIC_ASSUME`. This is needed to handle failed scenarios that
         // are replayed with a modified version of test driver (that use new `vm.assume`
         // cheatcodes).
-        if call_result.reverted && fail_on_revert && call_result.result.as_ref() != MAGIC_ASSUME {
+        if call_result.reverted
+            && options.fail_on_revert
+            && call_result.result.as_ref() != MAGIC_ASSUME
+        {
             // Candidate sequence fails test.
             // We don't have to apply remaining calls to check sequence.
-            return Ok((false, false));
+            return Ok((false, false, call_failure_reason(call_result, options.rd)));
         }
     }
 
     // Check the invariant for call sequence.
-    let (_, mut success) = call_invariant_function(&executor, test_address, calldata)?;
+    let (invariant_result, mut success) =
+        call_invariant_function(&executor, test_address, calldata)?;
+    if !success {
+        return Ok((false, true, call_failure_reason(invariant_result, options.rd)));
+    }
+
     // Check after invariant result if invariant is success and `afterInvariant` function is
     // declared.
-    if success && call_after_invariant {
-        (_, success) = call_after_invariant_function(&executor, test_address)?;
+    if success && options.call_after_invariant {
+        let (after_invariant_result, after_invariant_success) =
+            call_after_invariant_function(&executor, test_address)?;
+        success = after_invariant_success;
+        if !success {
+            return Ok((false, true, call_failure_reason(after_invariant_result, options.rd)));
+        }
     }
 
-    Ok((success, true))
+    Ok((success, true, None))
+}
+
+pub struct CheckSequenceOptions<'a> {
+    pub fail_on_revert: bool,
+    pub call_after_invariant: bool,
+    pub rd: Option<&'a RevertDecoder>,
+}
+
+fn call_failure_reason<FEN: FoundryEvmNetwork>(
+    call_result: RawCallResult<FEN>,
+    rd: Option<&RevertDecoder>,
+) -> Option<String> {
+    match call_result.into_evm_error(rd) {
+        EvmError::Execution(err) => Some(err.reason),
+        _ => None,
+    }
 }
 
 /// Shrinks a call sequence to the shortest sequence that still produces the target optimization

crates/forge/src/result.rs

@@ -681,6 +681,7 @@ impl TestResult {
         &mut self,
         replayed_entirely: bool,
         invariant_name: &String,
+        replay_reason: Option<String>,
         call_sequence: Vec<BaseCounterExample>,
     ) {
         self.kind = TestKind::Invariant {
@@ -692,11 +693,13 @@ impl TestResult {
             optimization_best_value: None,
         };
         self.status = TestStatus::Failure;
-        self.reason = if replayed_entirely {
-            Some(format!("{invariant_name} replay failure"))
-        } else {
-            Some(format!("{invariant_name} persisted failure revert"))
-        };
+        self.reason = replay_reason.or_else(|| {
+            if replayed_entirely {
+                Some(format!("{invariant_name} replay failure"))
+            } else {
+                Some(format!("{invariant_name} persisted failure revert"))
+            }
+        });
         self.counterexample = Some(CounterExample::Sequence(call_sequence.len(), call_sequence));
     }
 

crates/forge/src/runner.rs

@@ -23,7 +23,8 @@ use foundry_evm::{
         CallResult, EvmError, Executor, ITest, RawCallResult,
         fuzz::FuzzedExecutor,
         invariant::{
-            InvariantExecutor, InvariantFuzzError, check_sequence, replay_error, replay_run,
+            CheckSequenceOptions, InvariantExecutor, InvariantFuzzError, check_sequence,
+            replay_error, replay_run,
         },
     },
     fuzz::{
@@ -787,14 +788,17 @@ impl<'a> FunctionRunner<'a> {
                     }
                 })
                 .collect::<Vec<BasicTxDetails>>();
-            if let Ok((success, replayed_entirely)) = check_sequence(
+            if let Ok((success, replayed_entirely, replay_reason)) = check_sequence(
                 self.clone_executor(),
                 &txes,
                 (0..min(txes.len(), invariant_config.depth as usize)).collect(),
                 invariant_contract.address,
                 invariant_contract.invariant_function.selector().to_vec().into(),
-                invariant_config.fail_on_revert,
-                invariant_contract.call_after_invariant,
+                CheckSequenceOptions {
+                    fail_on_revert: invariant_config.fail_on_revert,
+                    call_after_invariant: invariant_contract.call_after_invariant,
+                    rd: Some(self.revert_decoder()),
+                },
             ) && !success
             {
                 let warn = format!(
@@ -845,6 +849,7 @@ impl<'a> FunctionRunner<'a> {
                 self.result.invariant_replay_fail(
                     replayed_entirely,
                     &invariant_contract.invariant_function.name,
+                    replay_reason,
                     call_sequence,
                 );
                 return self.result;

crates/forge/tests/cli/test_cmd/invariant/common.rs

@@ -1426,7 +1426,7 @@ Ran 1 test for test/InvariantShrinkBigSequence.t.sol:ShrinkBigSequenceTest
         cmd.assert_failure().stdout_eq(str![[r#"
 ...
 Ran 1 test for test/InvariantShrinkBigSequence.t.sol:ShrinkBigSequenceTest
-[FAIL: invariant_shrink_big_sequence replay failure]
+[FAIL: condition met]
 	[Sequence] (original: [..], shrunk: 77)
 ...
 "#]]);

crates/forge/tests/cli/test_cmd/invariant/mod.rs

@@ -474,7 +474,7 @@ Tip: Run `forge test --rerun` to retry only the 1 failed test
 ...
 Failing tests:
 Encountered 1 failing test in test/InvariantSequenceLenTest.t.sol:InvariantSequenceLenTest
-[FAIL: invariant_increment replay failure]
+[FAIL: invariant increment failure]
 	[Sequence] (original: 3, shrunk: 3)
 		sender=0x00000000000000000000000000000000000014aD addr=[src/Counter.sol:Counter]0x5615dEB798BB3E4dFa0139dFa1b3D433Cc23b72f calldata=increment() args=[]
 		sender=0x8ef7F804bAd9183981A366EA618d9D47D3124649 addr=[src/Counter.sol:Counter]0x5615dEB798BB3E4dFa0139dFa1b3D433Cc23b72f calldata=increment() args=[]
@@ -542,7 +542,7 @@ contract OwnableTest is Test {
     // Should replay failure if same test.
     cmd.assert_failure().stdout_eq(str![[r#"
 ...
-[FAIL: invariant_never_owner replay failure]
+[FAIL: never owner]
 ...
 "#]]);
 
@@ -582,6 +582,146 @@ Warning: Failure from "[..]/invariant/failures/OwnableTest/invariant_never_owner
 "#]]);
 });
 
+forgetest_init!(invariant_replay_preserves_fail_reason, |prj, cmd| {
+    prj.update_config(|config| {
+        config.invariant.runs = 1;
+        config.invariant.depth = 1;
+    });
+    prj.add_test(
+        "InvariantReplayFailReason.t.sol",
+        r#"
+import {Test} from "forge-std/Test.sol";
+
+contract InvariantReplayFailReason is Test {
+    function setUp() public {
+        targetContract(address(this));
+    }
+
+    function callTarget(uint256) external {}
+
+    function invariant_fail_reason() public {
+        fail();
+    }
+}
+   "#,
+    );
+
+    cmd.args(["test", "--mt", "invariant_fail_reason"]).assert_failure().stdout_eq(str![[r#"
+...
+[FAIL: failed to set up invariant testing environment: <empty revert data>][..]
+...
+"#]]);
+
+    // Replay should preserve failure reason instead of generic replay message.
+    cmd.assert_failure().stdout_eq(str![[r#"
+...
+[FAIL: failed to set up invariant testing environment: <empty revert data>][..]
+...
+"#]]);
+});
+
+forgetest_init!(invariant_replay_preserves_custom_error_reason, |prj, cmd| {
+    prj.update_config(|config| {
+        config.invariant.runs = 1;
+        config.invariant.depth = 1;
+        config.invariant.fail_on_revert = true;
+    });
+    prj.add_test(
+        "InvariantReplayCustomError.t.sol",
+        r#"
+import {Test} from "forge-std/Test.sol";
+
+contract CustomErrorTarget {
+    error InvariantCustomError(uint256, string);
+
+    function breakInvariant() external {
+        revert InvariantCustomError(111, "custom");
+    }
+}
+
+contract CustomErrorHandler is Test {
+    CustomErrorTarget target;
+
+    constructor() {
+        target = new CustomErrorTarget();
+    }
+
+    function callTarget() external {
+        target.breakInvariant();
+    }
+}
+
+contract InvariantReplayCustomError is Test {
+    CustomErrorHandler handler;
+
+    function setUp() public {
+        handler = new CustomErrorHandler();
+        targetContract(address(handler));
+    }
+
+    function invariant_custom_error_reason() public view {}
+}
+   "#,
+    );
+
+    cmd.args(["test", "--mt", "invariant_custom_error_reason"]).assert_failure().stdout_eq(str![[
+        r#"
+...
+[FAIL: [..]custom[..]][..]
+...
+"#
+    ]]);
+
+    // Replay should preserve custom error string too.
+    cmd.assert_failure().stdout_eq(str![[r#"
+...
+[FAIL: [..]custom[..]][..]
+...
+"#]]);
+});
+
+forgetest_init!(invariant_replay_preserves_invariant_custom_error_reason, |prj, cmd| {
+    prj.update_config(|config| {
+        config.invariant.runs = 1;
+        config.invariant.depth = 1;
+    });
+    prj.add_test(
+        "InvariantReplayInvariantCustomError.t.sol",
+        r#"
+import {Test} from "forge-std/Test.sol";
+
+contract InvariantReplayInvariantCustomError is Test {
+    error InvariantCustomError(uint256, string);
+
+    function setUp() public {
+        targetContract(address(this));
+    }
+
+    function touch(uint256) external {}
+
+    function invariant_custom_error_reason_from_invariant() public pure {
+        revert InvariantCustomError(222, "invariant custom");
+    }
+}
+   "#,
+    );
+
+    cmd.args(["test", "--mt", "invariant_custom_error_reason_from_invariant"])
+        .assert_failure()
+        .stdout_eq(str![[r#"
+...
+[FAIL: failed to set up invariant testing environment: InvariantCustomError(222, "invariant custom")][..]
+...
+"#]]);
+
+    // Replay should preserve invariant-level custom error string too.
+    cmd.assert_failure().stdout_eq(str![[r#"
+...
+[FAIL: failed to set up invariant testing environment: InvariantCustomError(222, "invariant custom")][..]
+...
+"#]]);
+});
+
 // <https://github.com/foundry-rs/foundry/issues/10253>
 forgetest_init!(invariant_test_target, |prj, cmd| {
     prj.update_config(|config| {