Use visit_assign to detect SSA locals.

3 files changed, 44 insertions, 7 deletions

diff --git a/compiler/rustc_mir_transform/src/ssa.rs b/compiler/rustc_mir_transform/src/ssa.rs
index 05a7b226f0c..a7b45366662 100644
--- a/compiler/rustc_mir_transform/src/ssa.rs
+++ b/compiler/rustc_mir_transform/src/ssa.rs
@@ -209,13 +209,6 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor {
         match ctxt {
             PlaceContext::MutatingUse(MutatingUseContext::Projection)
             | PlaceContext::NonMutatingUse(NonMutatingUseContext::Projection) => bug!(),
-            PlaceContext::MutatingUse(MutatingUseContext::Store) => {
-                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);
-                }
-            }
             // Anything can happen with raw pointers, so remove them.
             // We do not verify that all uses of the borrow dominate the assignment to `local`,
             // so we have to remove them too.
@@ -252,6 +245,19 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor {
             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))]
diff --git a/tests/mir-opt/copy-prop/partial_init.main.CopyProp.diff b/tests/mir-opt/copy-prop/partial_init.main.CopyProp.diff
new file mode 100644
index 00000000000..5866439055e
--- /dev/null
+++ b/tests/mir-opt/copy-prop/partial_init.main.CopyProp.diff
@@ -0,0 +1,13 @@
+- // MIR for `main` before CopyProp
++ // MIR for `main` after CopyProp
+  
+  fn main() -> () {
+      let mut _0: ();                      // return place in scope 0 at $DIR/partial_init.rs:+0:15: +0:15
+      let mut _1: (isize,);                // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+  
+      bb0: {
+          (_1.0: isize) = const 1_isize;   // scope 0 at $DIR/partial_init.rs:+4:13: +4:20
+          return;                          // scope 0 at $DIR/partial_init.rs:+5:13: +5:21
+      }
+  }
+  
diff --git a/tests/mir-opt/copy-prop/partial_init.rs b/tests/mir-opt/copy-prop/partial_init.rs
new file mode 100644
index 00000000000..f5ab9974f71
--- /dev/null
+++ b/tests/mir-opt/copy-prop/partial_init.rs
@@ -0,0 +1,18 @@
+// unit-test: CopyProp
+// Verify that we do not ICE on partial initializations.
+
+#![feature(custom_mir, core_intrinsics)]
+extern crate core;
+use core::intrinsics::mir::*;
+
+// EMIT_MIR partial_init.main.CopyProp.diff
+#[custom_mir(dialect = "runtime", phase = "post-cleanup")]
+pub fn main() {
+    mir! (
+        let x: (isize, );
+        {
+            x.0 = 1;
+            Return()
+        }
+    )
+}