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] 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 5922922d47b..3e8a0a2b7df 100644
--- a/compiler/rustc_const_eval/src/transform/validate.rs
+++ b/compiler/rustc_const_eval/src/transform/validate.rs
@@ -285,6 +285,12 @@ fn check_unwind_edge(&mut self, location: Location, unwind: UnwindAction) {
             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 @@ fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location
                 }
                 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 00000000000..9ef655cd1bb
--- /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()
+        }
+    )
+}