Rollup merge of rust-lang#65624 - wesleywiser:improve_simplify_locals, r=oli-obk

[mir-opt] Improve SimplifyLocals pass so it can remove unused consts

The `ConstProp` can cause many locals to be initialized to a constant
value and then never read from. `ConstProp` can also evaluate ZSTs into
constant values. Previously, many of these would be removed by other
parts of the MIR optimization pipeline. However, evaluating ZSTs
(especially `()`) into constant values defeated those parts of the
optimizer and so in a2e3ed5, I added a
hack to `ConstProp` that skips evaluating ZSTs to avoid that regression.

This commit changes `SimplifyLocals` so that it doesn't consider writes
of const values to a local to be a use of that local. In doing so,
`SimplifyLocals` is able to remove otherwise unused locals left behind
by other optimization passes (`ConstProp` in particular).

r? @oli-obk

Author: JohnTitor

src/librustc_mir/transform/const_prop.rs

@@ -431,7 +431,7 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
         place_layout: TyLayout<'tcx>,
         source_info: SourceInfo,
         place: &Place<'tcx>,
-    ) -> Option<Const<'tcx>> {
+    ) -> Option<()> {
         let span = source_info.span;
 
         let overflow_check = self.tcx.sess.overflow_checks();
@@ -540,20 +540,13 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
                 }
             }
 
-            // Work around: avoid extra unnecessary locals. FIXME(wesleywiser)
-            // Const eval will turn this into a `const Scalar(<ZST>)` that
-            // `SimplifyLocals` doesn't know it can remove.
-            Rvalue::Aggregate(_, operands) if operands.len() == 0 => {
-                return None;
-            }
-
             _ => { }
         }
 
         self.use_ecx(source_info, |this| {
             trace!("calling eval_rvalue_into_place(rvalue = {:?}, place = {:?})", rvalue, place);
             this.ecx.eval_rvalue_into_place(rvalue, place)?;
-            this.ecx.eval_place_to_op(place, Some(place_layout))
+            Ok(())
         })
     }
 
@@ -717,24 +710,23 @@ impl<'mir, 'tcx> MutVisitor<'tcx> for ConstPropagator<'mir, 'tcx> {
                     base: PlaceBase::Local(local),
                     projection: box [],
                 } = *place {
-                    if let Some(value) = self.const_prop(rval,
-                                                         place_layout,
-                                                         statement.source_info,
-                                                         place) {
-                        trace!("checking whether {:?} can be stored to {:?}", value, local);
+                    let source = statement.source_info;
+                    if let Some(()) = self.const_prop(rval, place_layout, source, place) {
                         if self.can_const_prop[local] {
-                            trace!("stored {:?} to {:?}", value, local);
-                            assert_eq!(self.get_const(local), Some(value));
+                            trace!("propagated into {:?}", local);
 
                             if self.should_const_prop() {
+                                let value =
+                                    self.get_const(local).expect("local was dead/uninitialized");
+                                trace!("replacing {:?} with {:?}", rval, value);
                                 self.replace_with_const(
                                     rval,
                                     value,
                                     statement.source_info,
                                 );
                             }
                         } else {
-                            trace!("can't propagate {:?} to {:?}", value, local);
+                            trace!("can't propagate into {:?}", local);
                             self.remove_const(local);
                         }
                     }

src/librustc_mir/transform/simplify.rs

@@ -31,7 +31,7 @@ use rustc_index::bit_set::BitSet;
 use rustc_index::vec::{Idx, IndexVec};
 use rustc::ty::TyCtxt;
 use rustc::mir::*;
-use rustc::mir::visit::{MutVisitor, Visitor, PlaceContext};
+use rustc::mir::visit::{MutVisitor, Visitor, PlaceContext, MutatingUseContext};
 use rustc::session::config::DebugInfo;
 use std::borrow::Cow;
 use crate::transform::{MirPass, MirSource};
@@ -293,23 +293,31 @@ pub fn remove_dead_blocks(body: &mut Body<'_>) {
 pub struct SimplifyLocals;
 
 impl<'tcx> MirPass<'tcx> for SimplifyLocals {
-    fn run_pass(&self, tcx: TyCtxt<'tcx>, _: MirSource<'tcx>, body: &mut Body<'tcx>) {
-        let mut marker = DeclMarker { locals: BitSet::new_empty(body.local_decls.len()) };
-        marker.visit_body(body);
-        // Return pointer and arguments are always live
-        marker.locals.insert(RETURN_PLACE);
-        for arg in body.args_iter() {
-            marker.locals.insert(arg);
-        }
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut Body<'tcx>) {
+        trace!("running SimplifyLocals on {:?}", source);
+        let locals = {
+            let mut marker = DeclMarker {
+                locals: BitSet::new_empty(body.local_decls.len()),
+                body,
+            };
+            marker.visit_body(body);
+            // Return pointer and arguments are always live
+            marker.locals.insert(RETURN_PLACE);
+            for arg in body.args_iter() {
+                marker.locals.insert(arg);
+            }
 
-        // We may need to keep dead user variables live for debuginfo.
-        if tcx.sess.opts.debuginfo == DebugInfo::Full {
-            for local in body.vars_iter() {
-                marker.locals.insert(local);
+            // We may need to keep dead user variables live for debuginfo.
+            if tcx.sess.opts.debuginfo == DebugInfo::Full {
+                for local in body.vars_iter() {
+                    marker.locals.insert(local);
+                }
             }
-        }
 
-        let map = make_local_map(&mut body.local_decls, marker.locals);
+            marker.locals
+        };
+
+        let map = make_local_map(&mut body.local_decls, locals);
         // Update references to all vars and tmps now
         LocalUpdater { map }.visit_body(body);
         body.local_decls.shrink_to_fit();
@@ -334,18 +342,35 @@ fn make_local_map<V>(
     map
 }
 
-struct DeclMarker {
+struct DeclMarker<'a, 'tcx> {
     pub locals: BitSet<Local>,
+    pub body: &'a Body<'tcx>,
 }
 
-impl<'tcx> Visitor<'tcx> for DeclMarker {
-    fn visit_local(&mut self, local: &Local, ctx: PlaceContext, _: Location) {
+impl<'a, 'tcx> Visitor<'tcx> for DeclMarker<'a, 'tcx> {
+    fn visit_local(&mut self, local: &Local, ctx: PlaceContext, location: Location) {
         // Ignore storage markers altogether, they get removed along with their otherwise unused
         // decls.
         // FIXME: Extend this to all non-uses.
-        if !ctx.is_storage_marker() {
-            self.locals.insert(*local);
+        if ctx.is_storage_marker() {
+            return;
         }
+
+        // Ignore stores of constants because `ConstProp` and `CopyProp` can remove uses of many
+        // of these locals. However, if the local is still needed, then it will be referenced in
+        // another place and we'll mark it as being used there.
+        if ctx == PlaceContext::MutatingUse(MutatingUseContext::Store) {
+            let stmt =
+                &self.body.basic_blocks()[location.block].statements[location.statement_index];
+            if let StatementKind::Assign(box (p, Rvalue::Use(Operand::Constant(c)))) = &stmt.kind {
+                if p.as_local().is_some() {
+                    trace!("skipping store of const value {:?} to {:?}", c, local);
+                    return;
+                }
+            }
+        }
+
+        self.locals.insert(*local);
     }
 }
 
@@ -357,9 +382,16 @@ impl<'tcx> MutVisitor<'tcx> for LocalUpdater {
     fn visit_basic_block_data(&mut self, block: BasicBlock, data: &mut BasicBlockData<'tcx>) {
         // Remove unnecessary StorageLive and StorageDead annotations.
         data.statements.retain(|stmt| {
-            match stmt.kind {
+            match &stmt.kind {
                 StatementKind::StorageLive(l) | StatementKind::StorageDead(l) => {
-                    self.map[l].is_some()
+                    self.map[*l].is_some()
+                }
+                StatementKind::Assign(box (place, _)) => {
+                    if let Some(local) = place.as_local() {
+                        self.map[local].is_some()
+                    } else {
+                        true
+                    }
                 }
                 _ => true
             }

src/test/incremental/hashes/for_loops.rs

@@ -25,7 +25,7 @@ pub fn change_loop_body() {
 }
 
 #[cfg(not(cfail1))]
-#[rustc_clean(cfg="cfail2", except="HirBody, mir_built, optimized_mir")]
+#[rustc_clean(cfg="cfail2", except="HirBody, mir_built")]
 #[rustc_clean(cfg="cfail3")]
 pub fn change_loop_body() {
     let mut _x = 0;

src/test/incremental/hashes/let_expressions.rs

@@ -22,7 +22,7 @@ pub fn change_name() {
 
 #[cfg(not(cfail1))]
 #[rustc_clean(cfg="cfail2",
-    except="HirBody,mir_built,optimized_mir")]
+    except="HirBody,mir_built")]
 #[rustc_clean(cfg="cfail3")]
 pub fn change_name() {
     let _y = 2u64;
@@ -86,7 +86,7 @@ pub fn change_mutability_of_slot() {
 
 #[cfg(not(cfail1))]
 #[rustc_clean(cfg="cfail2",
-    except="HirBody,typeck_tables_of,mir_built,optimized_mir")]
+    except="HirBody,typeck_tables_of,mir_built")]
 #[rustc_clean(cfg="cfail3")]
 pub fn change_mutability_of_slot() {
     let _x: u64 = 0;
@@ -182,7 +182,7 @@ pub fn add_initializer() {
 
 #[cfg(not(cfail1))]
 #[rustc_clean(cfg="cfail2",
-    except="HirBody,typeck_tables_of,mir_built,optimized_mir")]
+    except="HirBody,typeck_tables_of,mir_built")]
 #[rustc_clean(cfg="cfail3")]
 pub fn add_initializer() {
     let _x: i16 = 3i16;
@@ -198,7 +198,7 @@ pub fn change_initializer() {
 
 #[cfg(not(cfail1))]
 #[rustc_clean(cfg="cfail2",
-    except="HirBody,mir_built,optimized_mir")]
+    except="HirBody,mir_built")]
 #[rustc_clean(cfg="cfail3")]
 pub fn change_initializer() {
     let _x = 5u16;

src/test/incremental/hashes/loop_expressions.rs

@@ -25,7 +25,7 @@ pub fn change_loop_body() {
 }
 
 #[cfg(not(cfail1))]
-#[rustc_clean(cfg="cfail2", except="HirBody, mir_built, optimized_mir")]
+#[rustc_clean(cfg="cfail2", except="HirBody, mir_built")]
 #[rustc_clean(cfg="cfail3")]
 pub fn change_loop_body() {
     let mut _x = 0;

src/test/incremental/hashes/while_let_loops.rs

@@ -25,7 +25,7 @@ pub fn change_loop_body() {
 }
 
 #[cfg(not(cfail1))]
-#[rustc_clean(cfg="cfail2", except="HirBody, mir_built, optimized_mir")]
+#[rustc_clean(cfg="cfail2", except="HirBody, mir_built")]
 #[rustc_clean(cfg="cfail3")]
 pub fn change_loop_body() {
     let mut _x = 0;

src/test/incremental/hashes/while_loops.rs

@@ -25,7 +25,7 @@ pub fn change_loop_body() {
 }
 
 #[cfg(not(cfail1))]
-#[rustc_clean(cfg="cfail2", except="HirBody, mir_built, optimized_mir")]
+#[rustc_clean(cfg="cfail2", except="HirBody, mir_built")]
 #[rustc_clean(cfg="cfail3")]
 pub fn change_loop_body() {
     let mut _x = 0;

src/test/mir-opt/simplify-locals-removes-unused-consts.rs

@@ -0,0 +1,89 @@
+// compile-flags: -C overflow-checks=no
+
+fn use_zst(_: ((), ())) { }
+
+struct Temp {
+    x: u8
+}
+
+fn use_u8(_: u8) { }
+
+fn main() {
+    let ((), ()) = ((), ());
+    use_zst(((), ()));
+
+    use_u8((Temp { x : 40 }).x + 2);
+}
+
+// END RUST SOURCE
+
+// START rustc.main.SimplifyLocals.before.mir
+// let mut _0: ();
+// let mut _1: ((), ());
+// let mut _2: ();
+// let mut _3: ();
+// let _4: ();
+// let mut _5: ((), ());
+// let mut _6: ();
+// let mut _7: ();
+// let _8: ();
+// let mut _9: u8;
+// let mut _10: u8;
+// let mut _11: Temp;
+// scope 1 {
+// }
+// bb0: {
+//   StorageLive(_1);
+//   StorageLive(_2);
+//   _2 = const Scalar(<ZST>) : ();
+//   StorageLive(_3);
+//   _3 = const Scalar(<ZST>) : ();
+//   _1 = const Scalar(<ZST>) : ((), ());
+//   StorageDead(_3);
+//   StorageDead(_2);
+//   StorageDead(_1);
+//   StorageLive(_4);
+//   StorageLive(_6);
+//   _6 = const Scalar(<ZST>) : ();
+//   StorageLive(_7);
+//   _7 = const Scalar(<ZST>) : ();
+//   StorageDead(_7);
+//   StorageDead(_6);
+//   _4 = const use_zst(const Scalar(<ZST>) : ((), ())) -> bb1;
+// }
+// bb1: {
+//   StorageDead(_4);
+//   StorageLive(_8);
+//   StorageLive(_10);
+//   StorageLive(_11);
+//   _11 = const Scalar(0x28) : Temp;
+//   _10 = const 40u8;
+//   StorageDead(_10);
+//   _8 = const use_u8(const 42u8) -> bb2;
+// }
+// bb2: {
+//   StorageDead(_11);
+//   StorageDead(_8);
+//   return;
+// }
+// END rustc.main.SimplifyLocals.before.mir
+// START rustc.main.SimplifyLocals.after.mir
+// let mut _0: ();
+// let _1: ();
+// let _2: ();
+// scope 1 {
+// }
+// bb0: {
+//   StorageLive(_1);
+//   _1 = const use_zst(const Scalar(<ZST>) : ((), ())) -> bb1;
+// }
+// bb1: {
+//   StorageDead(_1);
+//   StorageLive(_2);
+//   _2 = const use_u8(const 42u8) -> bb2;
+// }
+// bb2: {
+//   StorageDead(_2);
+//   return;
+// }
+// END rustc.main.SimplifyLocals.after.mir

src/test/mir-opt/slice-drop-shim.rs

@@ -1,5 +1,7 @@
+// compile-flags: -Zmir-opt-level=0
+
 fn main() {
-    std::ptr::drop_in_place::<[String]> as unsafe fn(_);
+    let _fn = std::ptr::drop_in_place::<[String]> as unsafe fn(_);
 }
 
 // END RUST SOURCE