From bca761dc5d6d13fadc7a2f3c222ebdc86b3da4d3 Mon Sep 17 00:00:00 2001
From: Erik Desjardins <erikdesjardins@users.noreply.github.com>
Date: Tue, 16 Mar 2021 21:30:37 -0400
Subject: [PATCH] move to separate pass, cache layouts

---
 .../rustc_mir/src/transform/instcombine.rs    | 27 +---------
 compiler/rustc_mir/src/transform/mod.rs       |  2 +
 .../rustc_mir/src/transform/remove_zsts.rs    | 49 +++++++++++++++++++
 3 files changed, 52 insertions(+), 26 deletions(-)
 create mode 100644 compiler/rustc_mir/src/transform/remove_zsts.rs

diff --git a/compiler/rustc_mir/src/transform/instcombine.rs b/compiler/rustc_mir/src/transform/instcombine.rs
index 5ac2ff1fc9a..bad82fe893e 100644
--- a/compiler/rustc_mir/src/transform/instcombine.rs
+++ b/compiler/rustc_mir/src/transform/instcombine.rs
@@ -12,12 +12,10 @@ pub struct InstCombine;
 
 impl<'tcx> MirPass<'tcx> for InstCombine {
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-        let param_env = tcx.param_env(body.source.def_id());
         let (basic_blocks, local_decls) = body.basic_blocks_and_local_decls_mut();
-        let ctx = InstCombineContext { tcx, local_decls, param_env };
+        let ctx = InstCombineContext { tcx, local_decls };
         for block in basic_blocks.iter_mut() {
             for statement in block.statements.iter_mut() {
-                ctx.combine_zst(&statement.source_info, &mut statement.kind);
                 match statement.kind {
                     StatementKind::Assign(box (_place, ref mut rvalue)) => {
                         ctx.combine_bool_cmp(&statement.source_info, rvalue);
@@ -34,7 +32,6 @@ impl<'tcx> MirPass<'tcx> for InstCombine {
 struct InstCombineContext<'tcx, 'a> {
     tcx: TyCtxt<'tcx>,
     local_decls: &'a LocalDecls<'tcx>,
-    param_env: ty::ParamEnv<'tcx>,
 }
 
 impl<'tcx, 'a> InstCombineContext<'tcx, 'a> {
@@ -44,28 +41,6 @@ impl<'tcx, 'a> InstCombineContext<'tcx, 'a> {
         })
     }
 
-    /// Remove assignments to inhabited ZST places.
-    fn combine_zst(&self, source_info: &SourceInfo, kind: &mut StatementKind<'tcx>) {
-        match kind {
-            StatementKind::Assign(box (place, _)) => {
-                let place_ty = place.ty(self.local_decls, self.tcx).ty;
-                if let Ok(layout) = self.tcx.layout_of(self.param_env.and(place_ty)) {
-                    if layout.is_zst() && !layout.abi.is_uninhabited() {
-                        if self.tcx.consider_optimizing(|| {
-                            format!(
-                                "InstCombine ZST - Place: {:?} SourceInfo: {:?}",
-                                place, source_info
-                            )
-                        }) {
-                            *kind = StatementKind::Nop;
-                        }
-                    }
-                }
-            }
-            _ => {}
-        }
-    }
-
     /// Transform boolean comparisons into logical operations.
     fn combine_bool_cmp(&self, source_info: &SourceInfo, rvalue: &mut Rvalue<'tcx>) {
         match rvalue {
diff --git a/compiler/rustc_mir/src/transform/mod.rs b/compiler/rustc_mir/src/transform/mod.rs
index 13546442f66..c9187fdeced 100644
--- a/compiler/rustc_mir/src/transform/mod.rs
+++ b/compiler/rustc_mir/src/transform/mod.rs
@@ -44,6 +44,7 @@ pub mod promote_consts;
 pub mod remove_noop_landing_pads;
 pub mod remove_storage_markers;
 pub mod remove_unneeded_drops;
+pub mod remove_zsts;
 pub mod required_consts;
 pub mod rustc_peek;
 pub mod simplify;
@@ -494,6 +495,7 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     // The main optimizations that we do on MIR.
     let optimizations: &[&dyn MirPass<'tcx>] = &[
         &remove_storage_markers::RemoveStorageMarkers,
+        &remove_zsts::RemoveZsts,
         &const_goto::ConstGoto,
         &remove_unneeded_drops::RemoveUnneededDrops,
         &match_branches::MatchBranchSimplification,
diff --git a/compiler/rustc_mir/src/transform/remove_zsts.rs b/compiler/rustc_mir/src/transform/remove_zsts.rs
new file mode 100644
index 00000000000..838ffe9b596
--- /dev/null
+++ b/compiler/rustc_mir/src/transform/remove_zsts.rs
@@ -0,0 +1,49 @@
+//! Removes assignments to ZST places.
+
+use crate::transform::MirPass;
+use rustc_data_structures::fx::FxHashMap;
+use rustc_middle::mir::{Body, StatementKind};
+use rustc_middle::ty::TyCtxt;
+
+pub struct RemoveZsts;
+
+impl<'tcx> MirPass<'tcx> for RemoveZsts {
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        let param_env = tcx.param_env(body.source.def_id());
+
+        let (basic_blocks, local_decls) = body.basic_blocks_and_local_decls_mut();
+
+        let mut is_zst_cache = FxHashMap::default();
+
+        for block in basic_blocks.iter_mut() {
+            for statement in block.statements.iter_mut() {
+                match statement.kind {
+                    StatementKind::Assign(box (place, _)) => {
+                        let place_ty = place.ty(local_decls, tcx).ty;
+
+                        let is_inhabited_zst = *is_zst_cache.entry(place_ty).or_insert_with(|| {
+                            if let Ok(layout) = tcx.layout_of(param_env.and(place_ty)) {
+                                if layout.is_zst() && !layout.abi.is_uninhabited() {
+                                    return true;
+                                }
+                            }
+                            false
+                        });
+
+                        if is_inhabited_zst {
+                            if tcx.consider_optimizing(|| {
+                                format!(
+                                    "RemoveZsts - Place: {:?} SourceInfo: {:?}",
+                                    place, statement.source_info
+                                )
+                            }) {
+                                statement.make_nop();
+                            }
+                        }
+                    }
+                    _ => {}
+                }
+            }
+        }
+    }
+}