Check whether locals are too large instead of whether accesses into them are too large

Author: oli-obk

src/librustc_mir/transform/const_prop.rs

@@ -34,7 +34,9 @@ use crate::interpret::{
 };
 use crate::transform::{MirPass, MirSource};
 
-/// The maximum number of bytes that we'll allocate space for a return value.
+/// The maximum number of bytes that we'll allocate space for a local or the return value.
+/// Needed for #66397, because otherwise we eval into large places and that can cause OOM or just
+/// Severely regress performance.
 const MAX_ALLOC_LIMIT: u64 = 1024;
 
 /// Macro for machine-specific `InterpError` without allocation.
@@ -331,7 +333,10 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
         let param_env = tcx.param_env(def_id).with_reveal_all();
 
         let span = tcx.def_span(def_id);
-        let can_const_prop = CanConstProp::check(body);
+        // FIXME: `CanConstProp::check` computes the layout of all locals, return those layouts
+        // so we can write them to `ecx.frame_mut().locals.layout, reducing the duplication in
+        // `layout_of` query invocations.
+        let can_const_prop = CanConstProp::check(tcx, param_env, body);
         let mut only_propagate_inside_block_locals = BitSet::new_empty(can_const_prop.len());
         for (l, mode) in can_const_prop.iter_enumerated() {
             if *mode == ConstPropMode::OnlyInsideOwnBlock {
@@ -612,15 +617,9 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
     fn const_prop(
         &mut self,
         rvalue: &Rvalue<'tcx>,
-        place_layout: TyAndLayout<'tcx>,
         source_info: SourceInfo,
         place: Place<'tcx>,
     ) -> Option<()> {
-        // #66397: Don't try to eval into large places as that can cause an OOM
-        if place_layout.size >= Size::from_bytes(MAX_ALLOC_LIMIT) {
-            return None;
-        }
-
         // Perform any special handling for specific Rvalue types.
         // Generally, checks here fall into one of two categories:
         //   1. Additional checking to provide useful lints to the user
@@ -893,7 +892,11 @@ struct CanConstProp {
 
 impl CanConstProp {
     /// Returns true if `local` can be propagated
-    fn check(body: &Body<'_>) -> IndexVec<Local, ConstPropMode> {
+    fn check(
+        tcx: TyCtxt<'tcx>,
+        param_env: ParamEnv<'tcx>,
+        body: &Body<'tcx>,
+    ) -> IndexVec<Local, ConstPropMode> {
         let mut cpv = CanConstProp {
             can_const_prop: IndexVec::from_elem(ConstPropMode::FullConstProp, &body.local_decls),
             found_assignment: BitSet::new_empty(body.local_decls.len()),
@@ -903,6 +906,16 @@ impl CanConstProp {
             ),
         };
         for (local, val) in cpv.can_const_prop.iter_enumerated_mut() {
+            let ty = body.local_decls[local].ty;
+            match tcx.layout_of(param_env.and(ty)) {
+                Ok(layout) if layout.size < Size::from_bytes(MAX_ALLOC_LIMIT) => {}
+                // Either the layout fails to compute, then we can't use this local anyway
+                // or the local is too large, then we don't want to.
+                _ => {
+                    *val = ConstPropMode::NoPropagation;
+                    continue;
+                }
+            }
             // Cannot use args at all
             // Cannot use locals because if x < y { y - x } else { x - y } would
             //        lint for x != y
@@ -1018,61 +1031,52 @@ impl<'mir, 'tcx> MutVisitor<'tcx> for ConstPropagator<'mir, 'tcx> {
         let source_info = statement.source_info;
         self.source_info = Some(source_info);
         if let StatementKind::Assign(box (place, ref mut rval)) = statement.kind {
-            let place_ty: Ty<'tcx> = place.ty(&self.local_decls, self.tcx).ty;
-            if let Ok(place_layout) = self.tcx.layout_of(self.param_env.and(place_ty)) {
-                let can_const_prop = self.can_const_prop[place.local];
-                if let Some(()) = self.const_prop(rval, place_layout, source_info, place) {
-                    // This will return None if the above `const_prop` invocation only "wrote" a
-                    // type whose creation requires no write. E.g. a generator whose initial state
-                    // consists solely of uninitialized memory (so it doesn't capture any locals).
-                    if let Some(value) = self.get_const(place) {
-                        if self.should_const_prop(value) {
-                            trace!("replacing {:?} with {:?}", rval, value);
-                            self.replace_with_const(rval, value, source_info);
-                            if can_const_prop == ConstPropMode::FullConstProp
-                                || can_const_prop == ConstPropMode::OnlyInsideOwnBlock
-                            {
-                                trace!("propagated into {:?}", place);
-                            }
+            let can_const_prop = self.can_const_prop[place.local];
+            if let Some(()) = self.const_prop(rval, source_info, place) {
+                // This will return None if the above `const_prop` invocation only "wrote" a
+                // type whose creation requires no write. E.g. a generator whose initial state
+                // consists solely of uninitialized memory (so it doesn't capture any locals).
+                if let Some(value) = self.get_const(place) {
+                    if self.should_const_prop(value) {
+                        trace!("replacing {:?} with {:?}", rval, value);
+                        self.replace_with_const(rval, value, source_info);
+                        if can_const_prop == ConstPropMode::FullConstProp
+                            || can_const_prop == ConstPropMode::OnlyInsideOwnBlock
+                        {
+                            trace!("propagated into {:?}", place);
                         }
                     }
-                    match can_const_prop {
-                        ConstPropMode::OnlyInsideOwnBlock => {
-                            trace!(
-                                "found local restricted to its block. \
+                }
+                match can_const_prop {
+                    ConstPropMode::OnlyInsideOwnBlock => {
+                        trace!(
+                            "found local restricted to its block. \
                                 Will remove it from const-prop after block is finished. Local: {:?}",
-                                place.local
-                            );
-                        }
-                        ConstPropMode::OnlyPropagateInto | ConstPropMode::NoPropagation => {
-                            trace!("can't propagate into {:?}", place);
-                            if place.local != RETURN_PLACE {
-                                Self::remove_const(&mut self.ecx, place.local);
-                            }
+                            place.local
+                        );
+                    }
+                    ConstPropMode::OnlyPropagateInto | ConstPropMode::NoPropagation => {
+                        trace!("can't propagate into {:?}", place);
+                        if place.local != RETURN_PLACE {
+                            Self::remove_const(&mut self.ecx, place.local);
                         }
-                        ConstPropMode::FullConstProp => {}
                     }
-                } else {
-                    // Const prop failed, so erase the destination, ensuring that whatever happens
-                    // from here on, does not know about the previous value.
-                    // This is important in case we have
-                    // ```rust
-                    // let mut x = 42;
-                    // x = SOME_MUTABLE_STATIC;
-                    // // x must now be undefined
-                    // ```
-                    // FIXME: we overzealously erase the entire local, because that's easier to
-                    // implement.
-                    trace!(
-                        "propagation into {:?} failed.
-                        Nuking the entire site from orbit, it's the only way to be sure",
-                        place,
-                    );
-                    Self::remove_const(&mut self.ecx, place.local);
+                    ConstPropMode::FullConstProp => {}
                 }
             } else {
+                // Const prop failed, so erase the destination, ensuring that whatever happens
+                // from here on, does not know about the previous value.
+                // This is important in case we have
+                // ```rust
+                // let mut x = 42;
+                // x = SOME_MUTABLE_STATIC;
+                // // x must now be undefined
+                // ```
+                // FIXME: we overzealously erase the entire local, because that's easier to
+                // implement.
                 trace!(
-                    "cannot propagate into {:?}, because the type of the local is generic.",
+                    "propagation into {:?} failed.
+                        Nuking the entire site from orbit, it's the only way to be sure",
                     place,
                 );
                 Self::remove_const(&mut self.ecx, place.local);