From cef892ebab0cf01be3c29e9dd9ccb13f281c2371 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tomasz=20Mi=C4=85sko?= <tomasz.miasko@gmail.com>
Date: Sun, 19 Nov 2023 00:00:00 +0000
Subject: [PATCH 1/3] Validate there are no critical call edges in optimized
 MIR

---
 .../src/transform/validate.rs                 | 22 +++++++++++++
 tests/ui/mir/validate/critical-edge.rs        | 31 +++++++++++++++++++
 2 files changed, 53 insertions(+)
 create mode 100644 tests/ui/mir/validate/critical-edge.rs

diff --git a/compiler/rustc_const_eval/src/transform/validate.rs b/compiler/rustc_const_eval/src/transform/validate.rs
index 5922922d47b66..3e8a0a2b7df08 100644
--- a/compiler/rustc_const_eval/src/transform/validate.rs
+++ b/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.
diff --git a/tests/ui/mir/validate/critical-edge.rs b/tests/ui/mir/validate/critical-edge.rs
new file mode 100644
index 0000000000000..9ef655cd1bb4a
--- /dev/null
+++ b/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()
+        }
+    )
+}

From e0d6292ddbe6b44deec120b87952615c29465838 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tomasz=20Mi=C4=85sko?= <tomasz.miasko@gmail.com>
Date: Sun, 19 Nov 2023 00:00:00 +0000
Subject: [PATCH 2/3] Split critical call edges in coroutine drop shim

---
 compiler/rustc_mir_transform/src/coroutine.rs | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/compiler/rustc_mir_transform/src/coroutine.rs b/compiler/rustc_mir_transform/src/coroutine.rs
index dfafd8598306b..5890c0f2db389 100644
--- a/compiler/rustc_mir_transform/src/coroutine.rs
+++ b/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,
     );
 

From 329d015014a314f6131214d83d10b9d3b17d387d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tomasz=20Mi=C4=85sko?= <tomasz.miasko@gmail.com>
Date: Sun, 19 Nov 2023 00:00:00 +0000
Subject: [PATCH 3/3] Split critical call edges just before code generation

---
 compiler/rustc_mir_transform/src/shim.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/compiler/rustc_mir_transform/src/shim.rs b/compiler/rustc_mir_transform/src/shim.rs
index a9640146a5274..f24a2d07e4949 100644
--- a/compiler/rustc_mir_transform/src/shim.rs
+++ b/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)),
     );