Fix def-use check for call terminators

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, DefLocation::Argument);
+        analyzer.define(arg, DefLocation::Argument);
     }
 
     // If there exists a local definition that dominates all uses of that local,
@@ -74,7 +74,7 @@ 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: DefLocation) {
+    fn define(&mut self, local: mir::Local, location: DefLocation) {
         let kind = &mut self.locals[local];
         match *kind {
             LocalKind::ZST => {}
@@ -162,7 +162,7 @@ impl<'mir, 'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> Visitor<'tcx>
         debug!("visit_assign(place={:?}, rvalue={:?})", place, rvalue);
 
         if let Some(local) = place.as_local() {
-            self.assign(local, DefLocation::Body(location));
+            self.define(local, DefLocation::Assignment(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) {
@@ -183,9 +183,14 @@ impl<'mir, 'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> Visitor<'tcx>
 
     fn visit_local(&mut self, local: mir::Local, context: PlaceContext, location: Location) {
         match context {
-            PlaceContext::MutatingUse(MutatingUseContext::Call)
-            | PlaceContext::MutatingUse(MutatingUseContext::Yield) => {
-                self.assign(local, DefLocation::Body(location));
+            PlaceContext::MutatingUse(MutatingUseContext::Call) => {
+                let call = location.block;
+                let TerminatorKind::Call { target, .. } =
+                    self.fx.mir.basic_blocks[call].terminator().kind
+                else {
+                    bug!()
+                };
+                self.define(local, DefLocation::CallReturn { call, target });
             }
 
             PlaceContext::NonUse(_)
@@ -237,6 +242,8 @@ impl<'mir, 'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> Visitor<'tcx>
                     }
                 }
             }
+
+            PlaceContext::MutatingUse(MutatingUseContext::Yield) => bug!(),
         }
     }
 }

compiler/rustc_middle/src/mir/mod.rs

@@ -1611,14 +1611,29 @@ impl Location {
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 pub enum DefLocation {
     Argument,
-    Body(Location),
+    Assignment(Location),
+    CallReturn { call: BasicBlock, target: Option<BasicBlock> },
 }
 
 impl DefLocation {
     pub fn dominates(self, location: Location, dominators: &Dominators<BasicBlock>) -> bool {
         match self {
             DefLocation::Argument => true,
-            DefLocation::Body(def) => def.successor_within_block().dominates(location, dominators),
+            DefLocation::Assignment(def) => {
+                def.successor_within_block().dominates(location, dominators)
+            }
+            DefLocation::CallReturn { target: None, .. } => false,
+            DefLocation::CallReturn { call, target: Some(target) } => {
+                // The definition occurs on the call -> target edge. The definition dominates a use
+                // if and only if the edge is on all paths from the entry to the use.
+                //
+                // Note that a call terminator has only one edge that can reach the target, so when
+                // the call strongly dominates the target, all paths from the entry to the target
+                // go through the call -> target edge.
+                call != target
+                    && dominators.dominates(call, target)
+                    && dominators.dominates(target, location.block)
+            }
         }
     }
 }

compiler/rustc_mir_transform/src/ssa.rs

@@ -40,7 +40,8 @@ impl SsaLocals {
         let dominators = body.basic_blocks.dominators();
 
         let direct_uses = IndexVec::from_elem(0, &body.local_decls);
-        let mut visitor = SsaVisitor { assignments, assignment_order, dominators, direct_uses };
+        let mut visitor =
+            SsaVisitor { body, assignments, assignment_order, dominators, direct_uses };
 
         for local in body.args_iter() {
             visitor.assignments[local] = Set1::One(DefLocation::Argument);
@@ -110,7 +111,7 @@ impl SsaLocals {
         body: &'a Body<'tcx>,
     ) -> impl Iterator<Item = (Local, &'a Rvalue<'tcx>, Location)> + 'a {
         self.assignment_order.iter().filter_map(|&local| {
-            if let Set1::One(DefLocation::Body(loc)) = self.assignments[local] {
+            if let Set1::One(DefLocation::Assignment(loc)) = self.assignments[local] {
                 let stmt = body.stmt_at(loc).left()?;
                 // `loc` must point to a direct assignment to `local`.
                 let Some((target, rvalue)) = stmt.kind.as_assign() else { bug!() };
@@ -134,21 +135,21 @@ impl SsaLocals {
                     AssignedValue::Arg,
                     Location { block: START_BLOCK, statement_index: 0 },
                 ),
-                Set1::One(DefLocation::Body(loc)) => {
+                Set1::One(DefLocation::Assignment(loc)) => {
                     let bb = &mut basic_blocks[loc.block];
-                    let value = if loc.statement_index < bb.statements.len() {
-                        // `loc` must point to a direct assignment to `local`.
-                        let stmt = &mut bb.statements[loc.statement_index];
-                        let StatementKind::Assign(box (target, ref mut rvalue)) = stmt.kind else {
-                            bug!()
-                        };
-                        assert_eq!(target.as_local(), Some(local));
-                        AssignedValue::Rvalue(rvalue)
-                    } else {
-                        let term = bb.terminator_mut();
-                        AssignedValue::Terminator(&mut term.kind)
+                    // `loc` must point to a direct assignment to `local`.
+                    let stmt = &mut bb.statements[loc.statement_index];
+                    let StatementKind::Assign(box (target, ref mut rvalue)) = stmt.kind else {
+                        bug!()
                     };
-                    f(local, value, loc)
+                    assert_eq!(target.as_local(), Some(local));
+                    f(local, AssignedValue::Rvalue(rvalue), loc)
+                }
+                Set1::One(DefLocation::CallReturn { call, .. }) => {
+                    let bb = &mut basic_blocks[call];
+                    let loc = Location { block: call, statement_index: bb.statements.len() };
+                    let term = bb.terminator_mut();
+                    f(local, AssignedValue::Terminator(&mut term.kind), loc)
                 }
                 _ => {}
             }
@@ -201,14 +202,15 @@ impl SsaLocals {
     }
 }
 
-struct SsaVisitor<'a> {
+struct SsaVisitor<'tcx, 'a> {
+    body: &'a Body<'tcx>,
     dominators: &'a Dominators<BasicBlock>,
     assignments: IndexVec<Local, Set1<DefLocation>>,
     assignment_order: Vec<Local>,
     direct_uses: IndexVec<Local, u32>,
 }
 
-impl SsaVisitor<'_> {
+impl SsaVisitor<'_, '_> {
     fn check_dominates(&mut self, local: Local, loc: Location) {
         let set = &mut self.assignments[local];
         let assign_dominates = match *set {
@@ -224,7 +226,7 @@ impl SsaVisitor<'_> {
     }
 }
 
-impl<'tcx> Visitor<'tcx> for SsaVisitor<'_> {
+impl<'tcx> Visitor<'tcx> for SsaVisitor<'tcx, '_> {
     fn visit_local(&mut self, local: Local, ctxt: PlaceContext, loc: Location) {
         match ctxt {
             PlaceContext::MutatingUse(MutatingUseContext::Projection)
@@ -250,9 +252,18 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor<'_> {
 
     fn visit_place(&mut self, place: &Place<'tcx>, ctxt: PlaceContext, loc: Location) {
         let location = match ctxt {
-            PlaceContext::MutatingUse(
-                MutatingUseContext::Store | MutatingUseContext::Call | MutatingUseContext::Yield,
-            ) => Some(DefLocation::Body(loc)),
+            PlaceContext::MutatingUse(MutatingUseContext::Store) => {
+                Some(DefLocation::Assignment(loc))
+            }
+            PlaceContext::MutatingUse(MutatingUseContext::Call) => {
+                let call = loc.block;
+                let TerminatorKind::Call { target, .. } =
+                    self.body.basic_blocks[call].terminator().kind
+                else {
+                    bug!()
+                };
+                Some(DefLocation::CallReturn { call, target })
+            }
             _ => None,
         };
         if let Some(location) = location
@@ -359,7 +370,7 @@ impl StorageLiveLocals {
             for (statement_index, statement) in bbdata.statements.iter().enumerate() {
                 if let StatementKind::StorageLive(local) = statement.kind {
                     storage_live[local]
-                        .insert(DefLocation::Body(Location { block, statement_index }));
+                        .insert(DefLocation::Assignment(Location { block, statement_index }));
                 }
             }
         }

tests/ui/mir/ssa_call_ret.rs

@@ -0,0 +1,30 @@
+// Regression test for issue #117331, where variable `a` was misidentified as
+// being in SSA form (the definition occurs on the return edge only).
+//
+// edition:2021
+// compile-flags: --crate-type=lib
+// build-pass
+// needs-unwind
+#![feature(custom_mir, core_intrinsics)]
+use core::intrinsics::mir::*;
+
+#[custom_mir(dialect = "runtime", phase = "optimized")]
+pub fn f() -> u32 {
+    mir!(
+        let a: u32;
+        {
+            Call(a = g(), bb1, UnwindCleanup(bb2))
+        }
+        bb1 = {
+            RET = a;
+            Return()
+        }
+        bb2 (cleanup) = {
+            RET = a;
+            UnwindResume()
+        }
+    )
+}
+
+#[inline(never)]
+pub fn g() -> u32 { 0 }