fix: fork errors getting overriden by RevertDiagnostic inspector (#10954)

yash-atreya · web-flow · commit b74f48a2bb50 · 2025-07-08T15:25:32.000+05:30
* fix: fork revert diagnostic

* feat(`cheatcodes`): introduces `ForkRevertDiagnostic`
diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs
@@ -37,7 +37,7 @@ use foundry_common::{SELECTOR_LEN, TransactionMaybeSigned, evm::Breakpoints};
 use foundry_evm_core::{
     InspectorExt,
     abi::Vm::stopExpectSafeMemoryCall,
-    backend::{DatabaseError, DatabaseExt, RevertDiagnostic},
+    backend::{DatabaseError, DatabaseExt, ForkRevertDiagnostic},
     constants::{CHEATCODE_ADDRESS, HARDHAT_CONSOLE_ADDRESS, MAGIC_ASSUME},
     evm::{FoundryEvm, new_evm_with_existing_context},
 };
@@ -395,8 +395,8 @@ pub struct Cheatcodes {
     /// 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>,
+    /// Fork revert diagnostic information
+    pub fork_revert_diagnostic: Option<ForkRevertDiagnostic>,
 
     /// Recorded storage reads and writes
     pub accesses: RecordAccess,
@@ -666,6 +666,12 @@ impl Cheatcodes {
         call: &mut CallInputs,
         executor: &mut impl CheatcodesExecutor,
     ) -> Option<CallOutcome> {
+        // Initialize fork revert diagnostic for this call if we're in fork mode
+        if ecx.journaled_state.db().is_forked_mode() {
+            let revert_diag = self.fork_revert_diagnostic.get_or_insert_default();
+            revert_diag.current_call_target = Some(call.target_address);
+        }
+
         let gas = Gas::new(call.gas_limit);
         let curr_depth = ecx.journaled_state.depth();
 
@@ -1036,6 +1042,36 @@ impl Cheatcodes {
             None => false,
         }
     }
+
+    /// Checks if we should diagnose a fork revert for the current state
+    fn should_diagnose_fork_revert(&self, interpreter: &Interpreter, ecx: Ecx) -> Option<Address> {
+        // Only proceed if this is a REVERT opcode
+        if interpreter.bytecode.opcode() != op::REVERT {
+            return None;
+        }
+
+        // Only handle reverts with empty data (size = 0)
+        // Reverts with data already have meaningful error messages
+        let size = interpreter.stack.peek(1).ok()?;
+        if !size.is_zero() {
+            return None;
+        }
+
+        // Check if fork diagnostic tracking is enabled
+        // This is only set when we're in fork mode
+        let fork_diagnostic = self.fork_revert_diagnostic.as_ref()?;
+        let call_target = fork_diagnostic.current_call_target?;
+
+        // Only diagnose if the call target is different from the test contract
+        // We don't want to diagnose reverts within the test contract itself
+        if let TxKind::Call(test_contract) = ecx.tx.kind
+            && call_target != test_contract
+        {
+            return Some(call_target);
+        }
+
+        None
+    }
 }
 
 impl Inspector<EthEvmContext<&mut dyn DatabaseExt>> for Cheatcodes {
@@ -1063,6 +1099,20 @@ impl Inspector<EthEvmContext<&mut dyn DatabaseExt>> for Cheatcodes {
 
     #[inline]
     fn step(&mut self, interpreter: &mut Interpreter, ecx: Ecx) {
+        // Check for REVERT opcode to handle fork diagnostics before RevertDiagnostic
+        if let Some(call_target) = self.should_diagnose_fork_revert(interpreter, ecx) {
+            let journaled_state = ecx.journaled_state.clone();
+            if let Some(diagnostic) =
+                ecx.journaled_state.db().diagnose_revert(call_target, &journaled_state)
+            {
+                // Store the diagnostic
+                if let Some(ref mut fork_diagnostic) = self.fork_revert_diagnostic {
+                    fork_diagnostic.diagnostic = Some(diagnostic);
+                }
+                return;
+            }
+        }
+
         self.pc = interpreter.bytecode.pc();
 
         // `pauseGasMetering`: pause / resume interpreter gas.
@@ -1140,6 +1190,16 @@ impl Inspector<EthEvmContext<&mut dyn DatabaseExt>> for Cheatcodes {
     }
 
     fn call_end(&mut self, ecx: Ecx, call: &CallInputs, outcome: &mut CallOutcome) {
+        // Apply fork diagnostic if we have one
+        if let Some(mut fork_diagnostic) = self.fork_revert_diagnostic.take() {
+            if outcome.result.is_revert() {
+                if let Some(diagnostic) = fork_diagnostic.diagnostic.take() {
+                    outcome.result.output = Error::encode(diagnostic.to_error_msg(&self.labels));
+                    return;
+                }
+            }
+        }
+
         let cheatcode_call = call.target_address == CHEATCODE_ADDRESS
             || call.target_address == HARDHAT_CONSOLE_ADDRESS;
 
@@ -1411,33 +1471,7 @@ impl Inspector<EthEvmContext<&mut dyn DatabaseExt>> for Cheatcodes {
             self.expected_emits.clear()
         }
 
-        // this will ensure we don't have false positives when trying to diagnose reverts in fork
-        // mode
-        let diag = self.fork_revert_diagnostic.take();
-
-        // if there's a revert and a previous call was diagnosed as fork related revert then we can
-        // return a better error here
-        if outcome.result.is_revert()
-            && let Some(err) = diag
-        {
-            outcome.result.output = Error::encode(err.to_error_msg(&self.labels));
-            return;
-        }
-
-        // try to diagnose reverts in multi-fork mode where a call is made to an address that does
-        // not exist
-        if let TxKind::Call(test_contract) = ecx.tx.kind {
-            // if a call to a different contract than the original test contract returned with
-            // `Stop` we check if the contract actually exists on the active fork
-            if ecx.journaled_state.db().is_forked_mode()
-                && outcome.result.result == InstructionResult::Stop
-                && call.target_address != test_contract
-            {
-                let journaled_state = ecx.journaled_state.clone();
-                self.fork_revert_diagnostic =
-                    ecx.journaled_state.db().diagnose_revert(call.target_address, &journaled_state);
-            }
-        }
+        // Fork diagnostics are now handled in the step method to run before RevertDiagnostic
 
         // If the depth is 0, then this is the root call terminating
         if ecx.journaled_state.depth() == 0 {
diff --git a/crates/evm/core/src/backend/diagnostic.rs b/crates/evm/core/src/backend/diagnostic.rs
@@ -18,6 +18,15 @@ pub enum RevertDiagnostic {
     },
 }
 
+/// Contains fork-specific revert diagnostic information
+#[derive(Clone, Debug, Default)]
+pub struct ForkRevertDiagnostic {
+    /// The target address of the current call
+    pub current_call_target: Option<Address>,
+    /// The diagnostic information if available
+    pub diagnostic: Option<RevertDiagnostic>,
+}
+
 impl RevertDiagnostic {
     /// Converts the diagnostic to a readable error message
     pub fn to_error_msg(&self, labels: &AddressHashMap<String>) -> String {
diff --git a/crates/evm/core/src/backend/mod.rs b/crates/evm/core/src/backend/mod.rs
@@ -39,7 +39,7 @@ use std::{
 };
 
 mod diagnostic;
-pub use diagnostic::RevertDiagnostic;
+pub use diagnostic::{ForkRevertDiagnostic, RevertDiagnostic};
 
 mod error;
 pub use error::{BackendError, BackendResult, DatabaseError, DatabaseResult};
