From d5c1ad5ca1c014044a4bfd5c9c82fb277de304c3 Mon Sep 17 00:00:00 2001
From: Erik Desjardins <erikdesjardins@users.noreply.github.com>
Date: Wed, 24 Mar 2021 13:00:36 -0400
Subject: [PATCH] RemoveZsts: don't touch unions

---
 .../rustc_mir/src/transform/remove_zsts.rs    | 34 ++++++++++++++++++-
 ...ouch_unions.get_union.RemoveZsts.after.mir | 13 +++++++
 .../mir-opt/remove_zsts_dont_touch_unions.rs  | 19 +++++++++++
 3 files changed, 65 insertions(+), 1 deletion(-)
 create mode 100644 src/test/mir-opt/remove_zsts_dont_touch_unions.get_union.RemoveZsts.after.mir
 create mode 100644 src/test/mir-opt/remove_zsts_dont_touch_unions.rs

diff --git a/compiler/rustc_mir/src/transform/remove_zsts.rs b/compiler/rustc_mir/src/transform/remove_zsts.rs
index 144cedf593ede..70f7538dd57a8 100644
--- a/compiler/rustc_mir/src/transform/remove_zsts.rs
+++ b/compiler/rustc_mir/src/transform/remove_zsts.rs
@@ -1,7 +1,8 @@
 //! Removes assignments to ZST places.
 
 use crate::transform::MirPass;
-use rustc_middle::mir::{Body, StatementKind};
+use rustc_middle::mir::tcx::PlaceTy;
+use rustc_middle::mir::{Body, LocalDecls, Place, StatementKind};
 use rustc_middle::ty::{self, Ty, TyCtxt};
 
 pub struct RemoveZsts;
@@ -28,6 +29,9 @@ impl<'tcx> MirPass<'tcx> for RemoveZsts {
                         if !layout.is_zst() {
                             continue;
                         }
+                        if involves_a_union(place, local_decls, tcx) {
+                            continue;
+                        }
                         if tcx.consider_optimizing(|| {
                             format!(
                                 "RemoveZsts - Place: {:?} SourceInfo: {:?}",
@@ -55,3 +59,31 @@ fn maybe_zst(ty: Ty<'_>) -> bool {
         _ => false,
     }
 }
+
+/// Miri lazily allocates memory for locals on assignment,
+/// so we must preserve writes to unions and union fields,
+/// or it will ICE on reads of those fields.
+fn involves_a_union<'tcx>(
+    place: Place<'tcx>,
+    local_decls: &LocalDecls<'tcx>,
+    tcx: TyCtxt<'tcx>,
+) -> bool {
+    let mut place_ty = PlaceTy::from_ty(local_decls[place.local].ty);
+    if is_union(place_ty.ty) {
+        return true;
+    }
+    for elem in place.projection {
+        place_ty = place_ty.projection_ty(tcx, elem);
+        if is_union(place_ty.ty) {
+            return true;
+        }
+    }
+    return false;
+}
+
+fn is_union(ty: Ty<'_>) -> bool {
+    match ty.kind() {
+        ty::Adt(def, _) if def.is_union() => true,
+        _ => false,
+    }
+}
diff --git a/src/test/mir-opt/remove_zsts_dont_touch_unions.get_union.RemoveZsts.after.mir b/src/test/mir-opt/remove_zsts_dont_touch_unions.get_union.RemoveZsts.after.mir
new file mode 100644
index 0000000000000..33bd9eb9b190a
--- /dev/null
+++ b/src/test/mir-opt/remove_zsts_dont_touch_unions.get_union.RemoveZsts.after.mir
@@ -0,0 +1,13 @@
+// MIR for `get_union` after RemoveZsts
+
+fn get_union() -> Foo {
+    let mut _0: Foo;                     // return place in scope 0 at $DIR/remove_zsts_dont_touch_unions.rs:12:19: 12:22
+    let mut _1: ();                      // in scope 0 at $DIR/remove_zsts_dont_touch_unions.rs:13:14: 13:16
+
+    bb0: {
+        StorageLive(_1);                 // scope 0 at $DIR/remove_zsts_dont_touch_unions.rs:13:14: 13:16
+        (_0.0: ()) = move _1;            // scope 0 at $DIR/remove_zsts_dont_touch_unions.rs:13:5: 13:18
+        StorageDead(_1);                 // scope 0 at $DIR/remove_zsts_dont_touch_unions.rs:13:17: 13:18
+        return;                          // scope 0 at $DIR/remove_zsts_dont_touch_unions.rs:14:2: 14:2
+    }
+}
diff --git a/src/test/mir-opt/remove_zsts_dont_touch_unions.rs b/src/test/mir-opt/remove_zsts_dont_touch_unions.rs
new file mode 100644
index 0000000000000..7a6f86b808509
--- /dev/null
+++ b/src/test/mir-opt/remove_zsts_dont_touch_unions.rs
@@ -0,0 +1,19 @@
+// compile-flags: -Zmir-opt-level=3
+
+// Ensure RemoveZsts doesn't remove ZST assignments to union fields,
+// which causes problems in Miri.
+
+union Foo {
+    x: (),
+    y: u64,
+}
+
+// EMIT_MIR remove_zsts_dont_touch_unions.get_union.RemoveZsts.after.mir
+fn get_union() -> Foo {
+    Foo { x: () }
+}
+
+
+fn main() {
+    get_union();
+}