Also consider call and yield as MIR SSA.

compiler/rustc_mir_transform/src/gvn.rs

@@ -63,7 +63,7 @@ use rustc_middle::mir::*;
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_target::abi::{VariantIdx, FIRST_VARIANT};
 
-use crate::ssa::SsaLocals;
+use crate::ssa::{AssignedValue, SsaLocals};
 use crate::MirPass;
 
 pub struct GVN;
@@ -87,21 +87,28 @@ fn propagate_ssa<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     let dominators = body.basic_blocks.dominators().clone();
 
     let mut state = VnState::new(tcx, param_env, &ssa, &dominators, &body.local_decls);
-    for arg in body.args_iter() {
-        if ssa.is_ssa(arg) {
-            let value = state.new_opaque().unwrap();
-            state.assign(arg, value);
-        }
-    }
-
-    ssa.for_each_assignment_mut(&mut body.basic_blocks, |local, rvalue, location| {
-        let value = state.simplify_rvalue(rvalue, location).or_else(|| state.new_opaque()).unwrap();
-        // FIXME(#112651) `rvalue` may have a subtype to `local`. We can only mark `local` as
-        // reusable if we have an exact type match.
-        if state.local_decls[local].ty == rvalue.ty(state.local_decls, tcx) {
+    ssa.for_each_assignment_mut(
+        body.basic_blocks.as_mut_preserves_cfg(),
+        |local, value, location| {
+            let value = match value {
+                // We do not know anything of this assigned value.
+                AssignedValue::Arg | AssignedValue::Terminator(_) => None,
+                // Try to get some insight.
+                AssignedValue::Rvalue(rvalue) => {
+                    let value = state.simplify_rvalue(rvalue, location);
+                    // FIXME(#112651) `rvalue` may have a subtype to `local`. We can only mark `local` as
+                    // reusable if we have an exact type match.
+                    if state.local_decls[local].ty != rvalue.ty(state.local_decls, tcx) {
+                        return;
+                    }
+                    value
+                }
+            };
+            // `next_opaque` is `Some`, so `new_opaque` must return `Some`.
+            let value = value.or_else(|| state.new_opaque()).unwrap();
             state.assign(local, value);
-        }
-    });
+        },
+    );
 
     // Stop creating opaques during replacement as it is useless.
     state.next_opaque = None;

compiler/rustc_mir_transform/src/ssa.rs

@@ -49,9 +49,13 @@ impl SmallDominators<'_> {
         let assign_dominates = match *set {
             Set1::Empty | Set1::Many => false,
             Set1::One(LocationExtended::Arg) => true,
-            Set1::One(LocationExtended::Plain(assign)) => {
+            Set1::One(LocationExtended::Assign(assign)) => {
                 self.dominates(assign.successor_within_block(), loc)
             }
+            Set1::One(LocationExtended::Terminator(_, effect_bb)) => {
+                let effect_loc = Location { block: effect_bb, statement_index: 0 };
+                self.dominates(effect_loc, loc)
+            }
         };
         // We are visiting a use that is not dominated by an assignment.
         // Either there is a cycle involved, or we are reading for uninitialized local.
@@ -62,6 +66,12 @@ impl SmallDominators<'_> {
     }
 }
 
+pub enum AssignedValue<'a, 'tcx> {
+    Arg,
+    Rvalue(&'a mut Rvalue<'tcx>),
+    Terminator(&'a mut TerminatorKind<'tcx>),
+}
+
 impl SsaLocals {
     pub fn new<'tcx>(body: &Body<'tcx>) -> SsaLocals {
         let assignment_order = Vec::with_capacity(body.local_decls.len());
@@ -72,10 +82,12 @@ impl SsaLocals {
         let dominators = SmallDominators { inner: 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(LocationExtended::Arg);
+            visitor.assignment_order.push(local);
         }
 
         if body.basic_blocks.len() > 2 {
@@ -136,13 +148,16 @@ impl SsaLocals {
     ) -> bool {
         match self.assignments[local] {
             Set1::One(LocationExtended::Arg) => true,
-            Set1::One(LocationExtended::Plain(ass)) => {
+            Set1::One(LocationExtended::Assign(ass)) => {
                 if ass.block == location.block {
                     ass.statement_index < location.statement_index
                 } else {
                     dominators.dominates(ass.block, location.block)
                 }
             }
+            Set1::One(LocationExtended::Terminator(_, effect_bb)) => {
+                dominators.dominates(effect_bb, location.block)
+            }
             _ => false,
         }
     }
@@ -152,7 +167,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(LocationExtended::Plain(loc)) = self.assignments[local] {
+            if let Set1::One(LocationExtended::Assign(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!() };
@@ -166,18 +181,33 @@ impl SsaLocals {
 
     pub fn for_each_assignment_mut<'tcx>(
         &self,
-        basic_blocks: &mut BasicBlocks<'tcx>,
-        mut f: impl FnMut(Local, &mut Rvalue<'tcx>, Location),
+        basic_blocks: &mut IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
+        mut f: impl FnMut(Local, AssignedValue<'_, 'tcx>, Location),
     ) {
         for &local in &self.assignment_order {
-            if let Set1::One(LocationExtended::Plain(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];
-                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));
-                f(local, rvalue, loc)
+            match self.assignments[local] {
+                Set1::One(LocationExtended::Arg) => f(
+                    local,
+                    AssignedValue::Arg,
+                    Location { block: START_BLOCK, statement_index: 0 },
+                ),
+                Set1::One(LocationExtended::Assign(loc)) => {
+                    // `loc` must point to a direct assignment to `local`.
+                    let bb = &mut basic_blocks[loc.block];
+                    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));
+                    f(local, AssignedValue::Rvalue(rvalue), loc)
+                }
+                Set1::One(LocationExtended::Terminator(block, _)) => {
+                    let loc =
+                        Location { block, statement_index: basic_blocks[block].statements.len() };
+                    let term = basic_blocks[block].terminator_mut();
+                    f(local, AssignedValue::Terminator(&mut term.kind), loc)
+                }
+                _ => {}
             }
         }
     }
@@ -230,18 +260,24 @@ impl SsaLocals {
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 enum LocationExtended {
-    Plain(Location),
+    /// This local was assigned as a function parameter, before any execution.
     Arg,
+    /// `Location` points to the `Assign` statement.
+    Assign(Location),
+    /// The blocks are respectively the bb which contains the assignment,
+    /// and the bb in which the assignment effect is complete.
+    Terminator(BasicBlock, BasicBlock),
 }
 
-struct SsaVisitor<'a> {
+struct SsaVisitor<'a, 'tcx> {
+    body: &'a Body<'tcx>,
     dominators: SmallDominators<'a>,
     assignments: IndexVec<Local, Set1<LocationExtended>>,
     assignment_order: Vec<Local>,
     direct_uses: IndexVec<Local, u32>,
 }
 
-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)
@@ -266,34 +302,46 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor<'_> {
     }
 
     fn visit_place(&mut self, place: &Place<'tcx>, ctxt: PlaceContext, loc: Location) {
-        if place.projection.first() == Some(&PlaceElem::Deref) {
-            // Do not do anything for storage statements and debuginfo.
+        let location = match ctxt {
+            PlaceContext::MutatingUse(MutatingUseContext::Store) => {
+                Some(LocationExtended::Assign(loc))
+            }
+            PlaceContext::MutatingUse(MutatingUseContext::Call | MutatingUseContext::Yield) => {
+                // The assignment happens on the `loc -> target` edge. We need to ensure that
+                // this *edge* dominates all uses of the local. This is true if
+                // `loc` dominates `target` *and* `target` dominates all uses.
+                if let Some(target) = self.body.basic_blocks[loc.block].terminator().successors().next()
+                    && self.dominators.dominates(loc, Location { block: target, statement_index: 0 })
+                {
+                    Some(LocationExtended::Terminator(loc.block, target))
+                } else {
+                    None
+                }
+            }
+            _ => None,
+        };
+        if let Some(location) = location
+            && let Some(local) = place.as_local()
+        {
+            self.assignments[local].insert(location);
+            if let Set1::One(_) = self.assignments[local] {
+                // Only record if SSA-like, to avoid growing the vector needlessly.
+                self.assignment_order.push(local);
+            }
+        } else if place.projection.first() == Some(&PlaceElem::Deref) {
+            // Do not do anything for debuginfo.
             if ctxt.is_use() {
                 // Only change the context if it is a real use, not a "use" in debuginfo.
                 let new_ctxt = PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy);
 
                 self.visit_projection(place.as_ref(), new_ctxt, loc);
                 self.dominators.check_dominates(&mut self.assignments[place.local], loc);
             }
-            return;
         } else {
             self.visit_projection(place.as_ref(), ctxt, loc);
             self.visit_local(place.local, ctxt, loc);
         }
     }
-
-    fn visit_assign(&mut self, place: &Place<'tcx>, rvalue: &Rvalue<'tcx>, loc: Location) {
-        if let Some(local) = place.as_local() {
-            self.assignments[local].insert(LocationExtended::Plain(loc));
-            if let Set1::One(_) = self.assignments[local] {
-                // Only record if SSA-like, to avoid growing the vector needlessly.
-                self.assignment_order.push(local);
-            }
-        } else {
-            self.visit_place(place, PlaceContext::MutatingUse(MutatingUseContext::Store), loc);
-        }
-        self.visit_rvalue(rvalue, loc);
-    }
 }
 
 #[instrument(level = "trace", skip(ssa, body))]
@@ -376,7 +424,7 @@ impl StorageLiveLocals {
             for (statement_index, statement) in bbdata.statements.iter().enumerate() {
                 if let StatementKind::StorageLive(local) = statement.kind {
                     storage_live[local]
-                        .insert(LocationExtended::Plain(Location { block, statement_index }));
+                        .insert(LocationExtended::Assign(Location { block, statement_index }));
                 }
             }
         }

tests/mir-opt/copy-prop/calls.multiple_edges.CopyProp.diff

@@ -0,0 +1,21 @@
+- // MIR for `multiple_edges` before CopyProp
++ // MIR for `multiple_edges` after CopyProp
+
+  fn multiple_edges(_1: bool) -> u8 {
+      let mut _0: u8;
+      let mut _2: u8;
+
+      bb0: {
+          switchInt(_1) -> [1: bb1, otherwise: bb2];
+      }
+
+      bb1: {
+          _2 = dummy(const 13_u8) -> [return: bb2, unwind continue];
+      }
+
+      bb2: {
+          _0 = _2;
+          return;
+      }
+  }
+

tests/mir-opt/copy-prop/calls.nrvo.CopyProp.diff

@@ -0,0 +1,24 @@
+- // MIR for `nrvo` before CopyProp
++ // MIR for `nrvo` after CopyProp
+
+  fn nrvo() -> u8 {
+      let mut _0: u8;
+      let _1: u8;
+      scope 1 {
+-         debug y => _1;
++         debug y => _0;
+      }
+
+      bb0: {
+-         StorageLive(_1);
+-         _1 = dummy(const 5_u8) -> [return: bb1, unwind continue];
++         _0 = dummy(const 5_u8) -> [return: bb1, unwind continue];
+      }
+
+      bb1: {
+-         _0 = _1;
+-         StorageDead(_1);
+          return;
+      }
+  }
+

tests/mir-opt/copy-prop/calls.rs

@@ -0,0 +1,43 @@
+// Check that CopyProp does propagate return values of call terminators.
+// unit-test: CopyProp
+// needs-unwind
+
+#![feature(custom_mir, core_intrinsics)]
+use std::intrinsics::mir::*;
+
+#[inline(never)]
+fn dummy(x: u8) -> u8 {
+    x
+}
+
+// EMIT_MIR calls.nrvo.CopyProp.diff
+fn nrvo() -> u8 {
+    let y = dummy(5); // this should get NRVO
+    y
+}
+
+// EMIT_MIR calls.multiple_edges.CopyProp.diff
+#[custom_mir(dialect = "runtime", phase = "initial")]
+fn multiple_edges(t: bool) -> u8 {
+    mir! {
+        let x: u8;
+        {
+            match t { true => bbt, _ => ret }
+        }
+        bbt = {
+            Call(x = dummy(13), ret)
+        }
+        ret = {
+            // `x` is not assigned on the `bb0 -> ret` edge,
+            // so should not be marked as SSA for merging with `_0`.
+            RET = x;
+            Return()
+        }
+    }
+}
+
+fn main() {
+    // Make sure the function actually gets instantiated.
+    nrvo();
+    multiple_edges(false);
+}

tests/mir-opt/inline/inline_diverging.h.Inline.panic-abort.diff

@@ -8,11 +8,11 @@
 +     scope 1 (inlined call_twice::<!, fn() -> ! {sleep}>) {
 +         debug f => _2;
 +         let mut _3: &fn() -> ! {sleep};
-+         let mut _4: !;
++         let _4: !;
 +         let mut _5: &fn() -> ! {sleep};
-+         let mut _6: !;
 +         scope 2 {
 +             debug a => _4;
++             let _6: !;
 +             scope 3 {
 +                 debug b => _6;
 +             }

tests/mir-opt/inline/inline_diverging.h.Inline.panic-unwind.diff

@@ -10,10 +10,10 @@
 +         let mut _3: &fn() -> ! {sleep};
 +         let _4: !;
 +         let mut _5: &fn() -> ! {sleep};
-+         let mut _6: !;
 +         let mut _7: !;
 +         scope 2 {
 +             debug a => _4;
++             let _6: !;
 +             scope 3 {
 +                 debug b => _6;
 +             }