Rollup merge of #66216 - wesleywiser:const_prop_codegen_improvements, r=oli-obk

[mir-opt] Handle return place in ConstProp and improve SimplifyLocals pass

Temporarily rebased on top of #66074. The top 2 commits are new.

r? @oli-obk

src/librustc_mir/transform/const_prop.rs

@@ -9,7 +9,7 @@
 use rustc::mir::{
     AggregateKind, Constant, Location, Place, PlaceBase, Body, Operand, Rvalue, Local, UnOp,
     StatementKind, Statement, LocalKind, TerminatorKind, Terminator,  ClearCrossCrate, SourceInfo,
-    BinOp, SourceScope, SourceScopeLocalData, LocalDecl, BasicBlock,
+    BinOp, SourceScope, SourceScopeLocalData, LocalDecl, BasicBlock, RETURN_PLACE,
 };
 use rustc::mir::visit::{
     Visitor, PlaceContext, MutatingUseContext, MutVisitor, NonMutatingUseContext,
@@ -25,6 +25,7 @@
     LayoutOf, TyLayout, LayoutError, HasTyCtxt, TargetDataLayout, HasDataLayout,
 };
 
+use crate::rustc::ty::subst::Subst;
 use crate::interpret::{
     self, InterpCx, ScalarMaybeUndef, Immediate, OpTy,
     StackPopCleanup, LocalValue, LocalState, AllocId, Frame,
@@ -269,6 +270,7 @@ struct ConstPropagator<'mir, 'tcx> {
     param_env: ParamEnv<'tcx>,
     source_scope_local_data: ClearCrossCrate<IndexVec<SourceScope, SourceScopeLocalData>>,
     local_decls: IndexVec<Local, LocalDecl<'tcx>>,
+    ret: Option<OpTy<'tcx, ()>>,
 }
 
 impl<'mir, 'tcx> LayoutOf for ConstPropagator<'mir, 'tcx> {
@@ -308,11 +310,21 @@ fn new(
         let mut ecx = InterpCx::new(tcx.at(span), param_env, ConstPropMachine, ());
         let can_const_prop = CanConstProp::check(body);
 
+        let substs = &InternalSubsts::identity_for_item(tcx, def_id);
+
+        let ret =
+            ecx
+                .layout_of(body.return_ty().subst(tcx, substs))
+                .ok()
+                // Don't bother allocating memory for ZST types which have no values.
+                .filter(|ret_layout| !ret_layout.is_zst())
+                .map(|ret_layout| ecx.allocate(ret_layout, MemoryKind::Stack));
+
         ecx.push_stack_frame(
-            Instance::new(def_id, &InternalSubsts::identity_for_item(tcx, def_id)),
+            Instance::new(def_id, substs),
             span,
             dummy_body,
-            None,
+            ret.map(Into::into),
             StackPopCleanup::None {
                 cleanup: false,
             },
@@ -327,6 +339,7 @@ fn new(
             source_scope_local_data,
             //FIXME(wesleywiser) we can't steal this because `Visitor::super_visit_body()` needs it
             local_decls: body.local_decls.clone(),
+            ret: ret.map(Into::into),
         }
     }
 
@@ -335,6 +348,15 @@ fn release_stolen_data(self) -> ClearCrossCrate<IndexVec<SourceScope, SourceScop
     }
 
     fn get_const(&self, local: Local) -> Option<Const<'tcx>> {
+        if local == RETURN_PLACE {
+            // Try to read the return place as an immediate so that if it is representable as a
+            // scalar, we can handle it as such, but otherwise, just return the value as is.
+            return match self.ret.map(|ret| self.ecx.try_read_immediate(ret)) {
+                Some(Ok(Ok(imm))) => Some(imm.into()),
+                _ => self.ret,
+            };
+        }
+
         self.ecx.access_local(self.ecx.frame(), local, None).ok()
     }
 
@@ -643,7 +665,8 @@ fn check(body: &Body<'_>) -> IndexVec<Local, bool> {
             //        lint for x != y
             // FIXME(oli-obk): lint variables until they are used in a condition
             // FIXME(oli-obk): lint if return value is constant
-            *val = body.local_kind(local) == LocalKind::Temp;
+            let local_kind = body.local_kind(local);
+            *val = local_kind == LocalKind::Temp || local_kind == LocalKind::ReturnPointer;
 
             if !*val {
                 trace!("local {:?} can't be propagated because it's not a temporary", local);
@@ -731,7 +754,9 @@ fn visit_statement(
                             }
                         } else {
                             trace!("can't propagate into {:?}", local);
-                            self.remove_const(local);
+                            if local != RETURN_PLACE {
+                                self.remove_const(local);
+                            }
                         }
                     }
                 }

src/librustc_mir/transform/simplify.rs

@@ -359,13 +359,20 @@ fn visit_local(&mut self, local: &Local, ctx: PlaceContext, location: Location)
         // 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;
+        if ctx == PlaceContext::MutatingUse(MutatingUseContext::Store) ||
+           ctx == PlaceContext::MutatingUse(MutatingUseContext::Projection) {
+            let block = &self.body.basic_blocks()[location.block];
+            if location.statement_index != block.statements.len() {
+                let stmt =
+                    &block.statements[location.statement_index];
+
+                if let StatementKind::Assign(
+                    box (p, Rvalue::Use(Operand::Constant(c)))
+                ) = &stmt.kind {
+                    if !p.is_indirect() {
+                        trace!("skipping store of const value {:?} to {:?}", c, p);
+                        return;
+                    }
                 }
             }
         }
@@ -392,7 +399,7 @@ fn visit_basic_block_data(&mut self, block: BasicBlock, data: &mut BasicBlockDat
                     self.map[*l].is_some()
                 }
                 StatementKind::Assign(box (place, _)) => {
-                    if let Some(local) = place.as_local() {
+                    if let PlaceBase::Local(local) = place.base {
                         self.map[local].is_some()
                     } else {
                         true

src/test/incremental/hashes/struct_constructors.rs

@@ -152,7 +152,7 @@ pub fn change_constructor_path_regular_struct() {
 }
 
 #[cfg(not(cfail1))]
-#[rustc_clean(cfg="cfail2", except="HirBody,optimized_mir,mir_built,typeck_tables_of")]
+#[rustc_clean(cfg="cfail2", except="HirBody,mir_built,typeck_tables_of")]
 #[rustc_clean(cfg="cfail3")]
 pub fn change_constructor_path_regular_struct() {
     let _ = RegularStruct2 {
@@ -213,7 +213,7 @@ pub fn change_constructor_path_tuple_struct() {
 }
 
 #[cfg(not(cfail1))]
-#[rustc_clean(cfg="cfail2", except="HirBody,optimized_mir,mir_built,typeck_tables_of")]
+#[rustc_clean(cfg="cfail2", except="HirBody,mir_built,typeck_tables_of")]
 #[rustc_clean(cfg="cfail3")]
 pub fn change_constructor_path_tuple_struct() {
     let _ = TupleStruct2(0, 1, 2);

src/test/mir-opt/const_prop/return_place.rs

@@ -0,0 +1,54 @@
+// compile-flags: -C overflow-checks=on
+
+fn add() -> u32 {
+    2 + 2
+}
+
+fn main() {
+    add();
+}
+
+// END RUST SOURCE
+// START rustc.add.ConstProp.before.mir
+// fn add() -> u32 {
+//     let mut _0: u32;
+//     let mut _1: (u32, bool);
+//     bb0: {
+//         _1 = CheckedAdd(const 2u32, const 2u32);
+//         assert(!move (_1.1: bool), "attempt to add with overflow") -> bb1;
+//     }
+//     bb1: {
+//         _0 = move (_1.0: u32);
+//         return;
+//     }
+//     bb2 (cleanup): {
+//         resume;
+//     }
+// }
+// END rustc.add.ConstProp.before.mir
+// START rustc.add.ConstProp.after.mir
+// fn add() -> u32 {
+//     let mut _0: u32;
+//     let mut _1: (u32, bool);
+//     bb0: {
+//         _1 = (const 4u32, const false);
+//         assert(!const false, "attempt to add with overflow") -> bb1;
+//     }
+//     bb1: {
+//         _0 = const 4u32;
+//         return;
+//     }
+//     bb2 (cleanup): {
+//         resume;
+//     }
+// }
+// END rustc.add.ConstProp.after.mir
+// START rustc.add.PreCodegen.before.mir
+// fn add() -> u32 {
+//     let mut _0: u32;
+//     bb0: {
+//         _0 = const 4u32;
+//         return;
+//     }
+// }
+// END rustc.add.PreCodegen.before.mir