Validate there are no critical call edges in optimized MIR

compiler/rustc_const_eval/src/transform/validate.rs

@@ -285,6 +285,12 @@ impl<'a, 'tcx> CfgChecker<'a, 'tcx> {
             UnwindAction::Unreachable | UnwindAction::Terminate(UnwindTerminateReason::Abi) => (),
         }
     }
+
+    fn is_critical_call_edge(&self, target: Option<BasicBlock>, unwind: UnwindAction) -> bool {
+        let Some(target) = target else { return false };
+        matches!(unwind, UnwindAction::Cleanup(_) | UnwindAction::Terminate(_))
+            && self.body.basic_blocks.predecessors()[target].len() > 1
+    }
 }
 
 impl<'a, 'tcx> Visitor<'tcx> for CfgChecker<'a, 'tcx> {
@@ -425,6 +431,22 @@ impl<'a, 'tcx> Visitor<'tcx> for CfgChecker<'a, 'tcx> {
                 }
                 self.check_unwind_edge(location, *unwind);
 
+                // The code generation assumes that there are no critical call edges. The assumption
+                // is used to simplify inserting code that should be executed along the return edge
+                // from the call. FIXME(tmiasko): Since this is a strictly code generation concern,
+                // the code generation should be responsible for handling it.
+                if self.mir_phase >= MirPhase::Runtime(RuntimePhase::Optimized)
+                    && self.is_critical_call_edge(*target, *unwind)
+                {
+                    self.fail(
+                        location,
+                        format!(
+                            "encountered critical edge in `Call` terminator {:?}",
+                            terminator.kind,
+                        ),
+                    );
+                }
+
                 // The call destination place and Operand::Move place used as an argument might be
                 // passed by a reference to the callee. Consequently they must be non-overlapping
                 // and cannot be packed. Currently this simply checks for duplicate places.

compiler/rustc_mir_transform/src/coroutine.rs

@@ -51,6 +51,7 @@
 //! Otherwise it drops all the values in scope at the last suspension point.
 
 use crate::abort_unwinding_calls;
+use crate::add_call_guards;
 use crate::deref_separator::deref_finder;
 use crate::errors;
 use crate::pass_manager as pm;
@@ -1162,7 +1163,7 @@ fn create_coroutine_drop_shim<'tcx>(
     pm::run_passes_no_validate(
         tcx,
         &mut body,
-        &[&abort_unwinding_calls::AbortUnwindingCalls],
+        &[&abort_unwinding_calls::AbortUnwindingCalls, &add_call_guards::CriticalCallEdges],
         None,
     );
 

compiler/rustc_mir_transform/src/shim.rs

@@ -111,8 +111,8 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> Body<'
             &deref_separator::Derefer,
             &remove_noop_landing_pads::RemoveNoopLandingPads,
             &simplify::SimplifyCfg::MakeShim,
-            &add_call_guards::CriticalCallEdges,
             &abort_unwinding_calls::AbortUnwindingCalls,
+            &add_call_guards::CriticalCallEdges,
         ],
         Some(MirPhase::Runtime(RuntimePhase::Optimized)),
     );

tests/ui/mir/validate/critical-edge.rs

@@ -0,0 +1,31 @@
+// Optimized MIR shouldn't have critical call edges
+//
+// build-fail
+// edition: 2021
+// compile-flags: --crate-type=lib
+// failure-status: 101
+// dont-check-compiler-stderr
+// error-pattern: encountered critical edge in `Call` terminator
+#![feature(custom_mir, core_intrinsics)]
+use core::intrinsics::mir::*;
+
+#[custom_mir(dialect = "runtime", phase = "optimized")]
+#[inline(always)]
+pub fn f(a: u32) -> u32 {
+    mir!(
+        {
+            match a {
+                0 => bb1,
+                _ => bb2,
+            }
+        }
+        bb1 = {
+            Call(RET = f(1), bb2, UnwindTerminate(ReasonAbi))
+        }
+
+        bb2 = {
+            RET = 2;
+            Return()
+        }
+    )
+}