Check dominance for the full edge.

Author: cjgillot

compiler/rustc_mir_transform/src/ssa.rs

@@ -246,8 +246,16 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor<'_, 'tcx> {
                 Some(LocationExtended::Assign(loc))
             }
             PlaceContext::MutatingUse(MutatingUseContext::Call | MutatingUseContext::Yield) => {
-                let target = self.body.basic_blocks[loc.block].terminator().successors().next();
-                target.map(|target| LocationExtended::Terminator(loc.block, target))
+                // 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,
         };

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);
+}