...

Author: tmiasko

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 {
+                    unreachable!()
+                };
+                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) => unreachable!(),
         }
     }
 }

compiler/rustc_middle/src/mir/mod.rs

@@ -1567,14 +1567,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

@@ -104,7 +104,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] {
                 // `loc` must point to a direct assignment to `local`.
                 let Either::Left(stmt) = body.stmt_at(loc) else { bug!() };
                 let Some((target, rvalue)) = stmt.kind.as_assign() else { bug!() };
@@ -122,7 +122,7 @@ impl SsaLocals {
         mut f: impl FnMut(Local, &mut Rvalue<'tcx>, Location),
     ) {
         for &local in &self.assignment_order {
-            if let Set1::One(DefLocation::Body(loc)) = self.assignments[local] {
+            if let Set1::One(DefLocation::Assignment(loc)) = self.assignments[local] {
                 // `loc` must point to a direct assignment to `local`.
                 let bbs = basic_blocks.as_mut_preserves_cfg();
                 let bb = &mut bbs[loc.block];
@@ -246,7 +246,7 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor<'_> {
 
     fn visit_assign(&mut self, place: &Place<'tcx>, rvalue: &Rvalue<'tcx>, loc: Location) {
         if let Some(local) = place.as_local() {
-            self.assignments[local].insert(DefLocation::Body(loc));
+            self.assignments[local].insert(DefLocation::Assignment(loc));
             if let Set1::One(_) = self.assignments[local] {
                 // Only record if SSA-like, to avoid growing the vector needlessly.
                 self.assignment_order.push(local);
@@ -338,7 +338,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 }));
                 }
             }
         }