From e489971902c814e5adb7041d1bda230a1acb4fd7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tomasz=20Mi=C4=85sko?= <tomasz.miasko@gmail.com>
Date: Sun, 22 Jan 2023 00:00:00 +0000
Subject: [PATCH] Fix def-use dominance check

A definition does not dominate a use in the same statement. For example
in MIR generated for compound assignment x += a (when overflow checks
are disabled).
---
 compiler/rustc_codegen_ssa/src/mir/analyze.rs | 33 ++++++++++++-------
 tests/ui/mir/mir_codegen_ssa.rs               | 19 +++++++++++
 2 files changed, 41 insertions(+), 11 deletions(-)
 create mode 100644 tests/ui/mir/mir_codegen_ssa.rs

diff --git a/compiler/rustc_codegen_ssa/src/mir/analyze.rs b/compiler/rustc_codegen_ssa/src/mir/analyze.rs
index dd1ac2c74ae..95aad10fdb0 100644
--- a/compiler/rustc_codegen_ssa/src/mir/analyze.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/analyze.rs
@@ -36,7 +36,7 @@ pub fn non_ssa_locals<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
 
     // Arguments get assigned to by means of the function being called
     for arg in mir.args_iter() {
-        analyzer.assign(arg, mir::START_BLOCK.start_location());
+        analyzer.assign(arg, DefLocation::Argument);
     }
 
     // If there exists a local definition that dominates all uses of that local,
@@ -64,7 +64,22 @@ enum LocalKind {
     /// A scalar or a scalar pair local that is neither defined nor used.
     Unused,
     /// A scalar or a scalar pair local with a single definition that dominates all uses.
-    SSA(mir::Location),
+    SSA(DefLocation),
+}
+
+#[derive(Copy, Clone, PartialEq, Eq)]
+enum DefLocation {
+    Argument,
+    Body(Location),
+}
+
+impl DefLocation {
+    fn dominates(self, location: Location, dominators: &Dominators<mir::BasicBlock>) -> bool {
+        match self {
+            DefLocation::Argument => true,
+            DefLocation::Body(def) => def.successor_within_block().dominates(location, dominators),
+        }
+    }
 }
 
 struct LocalAnalyzer<'mir, 'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> {
@@ -74,17 +89,13 @@ struct LocalAnalyzer<'mir, 'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> {
 }
 
 impl<'mir, 'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> LocalAnalyzer<'mir, 'a, 'tcx, Bx> {
-    fn assign(&mut self, local: mir::Local, location: Location) {
+    fn assign(&mut self, local: mir::Local, location: DefLocation) {
         let kind = &mut self.locals[local];
         match *kind {
             LocalKind::ZST => {}
             LocalKind::Memory => {}
-            LocalKind::Unused => {
-                *kind = LocalKind::SSA(location);
-            }
-            LocalKind::SSA(_) => {
-                *kind = LocalKind::Memory;
-            }
+            LocalKind::Unused => *kind = LocalKind::SSA(location),
+            LocalKind::SSA(_) => *kind = LocalKind::Memory,
         }
     }
 
@@ -166,7 +177,7 @@ fn visit_assign(
         debug!("visit_assign(place={:?}, rvalue={:?})", place, rvalue);
 
         if let Some(local) = place.as_local() {
-            self.assign(local, location);
+            self.assign(local, DefLocation::Body(location));
             if self.locals[local] != LocalKind::Memory {
                 let decl_span = self.fx.mir.local_decls[local].source_info.span;
                 if !self.fx.rvalue_creates_operand(rvalue, decl_span) {
@@ -189,7 +200,7 @@ fn visit_local(&mut self, local: mir::Local, context: PlaceContext, location: Lo
         match context {
             PlaceContext::MutatingUse(MutatingUseContext::Call)
             | PlaceContext::MutatingUse(MutatingUseContext::Yield) => {
-                self.assign(local, location);
+                self.assign(local, DefLocation::Body(location));
             }
 
             PlaceContext::NonUse(_) | PlaceContext::MutatingUse(MutatingUseContext::Retag) => {}
diff --git a/tests/ui/mir/mir_codegen_ssa.rs b/tests/ui/mir/mir_codegen_ssa.rs
new file mode 100644
index 00000000000..5e2f10cefe9
--- /dev/null
+++ b/tests/ui/mir/mir_codegen_ssa.rs
@@ -0,0 +1,19 @@
+// build-pass
+// compile-flags: --crate-type=lib
+#![feature(custom_mir, core_intrinsics)]
+use std::intrinsics::mir::*;
+
+#[custom_mir(dialect = "runtime", phase = "optimized")]
+pub fn f(a: u32) -> u32 {
+    mir!(
+        let x: u32;
+        {
+            // Previously code generation failed with ICE "use of .. before def ..." because the
+            // definition of x was incorrectly identified as dominating the use of x located in the
+            // same statement:
+            x = x + a;
+            RET = x;
+            Return()
+        }
+    )
+}