Auto merge of #115797 - cjgillot:const-prop-noclone, r=oli-obk

Do not clone the Body for ConstProp

~Based on https://github.com/rust-lang/rust/pull/115748 for the `POST_MONO_CHECKS` flag.~

compiler/rustc_const_eval/src/interpret/eval_context.rs

@@ -749,10 +749,12 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         self.stack_mut().push(frame);
 
         // Make sure all the constants required by this frame evaluate successfully (post-monomorphization check).
-        for ct in &body.required_consts {
+        if M::POST_MONO_CHECKS {
-            let span = ct.span;
+            for ct in &body.required_consts {
-            let ct = self.subst_from_current_frame_and_normalize_erasing_regions(ct.literal)?;
+                let span = ct.span;
-            self.eval_mir_constant(&ct, Some(span), None)?;
+                let ct = self.subst_from_current_frame_and_normalize_erasing_regions(ct.literal)?;
+                self.eval_mir_constant(&ct, Some(span), None)?;
+            }
         }
 
         // done

compiler/rustc_const_eval/src/interpret/machine.rs

@@ -130,6 +130,9 @@ pub trait Machine<'mir, 'tcx: 'mir>: Sized {
     /// Should the machine panic on allocation failures?
     const PANIC_ON_ALLOC_FAIL: bool;
 
+    /// Should post-monomorphization checks be run when a stack frame is pushed?
+    const POST_MONO_CHECKS: bool = true;
+
     /// Whether memory accesses should be alignment-checked.
     fn enforce_alignment(ecx: &InterpCx<'mir, 'tcx, Self>) -> CheckAlignment;
 

compiler/rustc_mir_transform/src/const_prop.rs

@@ -15,10 +15,11 @@ use rustc_middle::mir::visit::{
 use rustc_middle::mir::*;
 use rustc_middle::ty::layout::{LayoutError, LayoutOf, LayoutOfHelpers, TyAndLayout};
 use rustc_middle::ty::{self, GenericArgs, Instance, ParamEnv, Ty, TyCtxt, TypeVisitableExt};
-use rustc_span::{def_id::DefId, Span, DUMMY_SP};
+use rustc_span::{def_id::DefId, Span};
 use rustc_target::abi::{self, Align, HasDataLayout, Size, TargetDataLayout};
 use rustc_target::spec::abi::Abi as CallAbi;
 
+use crate::dataflow_const_prop::Patch;
 use crate::MirPass;
 use rustc_const_eval::interpret::{
     self, compile_time_machine, AllocId, ConstAllocation, ConstValue, FnArg, Frame, ImmTy,
@@ -83,9 +84,9 @@ impl<'tcx> MirPass<'tcx> for ConstProp {
             return;
         }
 
-        let is_generator = tcx.type_of(def_id.to_def_id()).instantiate_identity().is_generator();
         // FIXME(welseywiser) const prop doesn't work on generators because of query cycles
         // computing their layout.
+        let is_generator = def_kind == DefKind::Generator;
         if is_generator {
             trace!("ConstProp skipped for generator {:?}", def_id);
             return;
@@ -93,33 +94,22 @@ impl<'tcx> MirPass<'tcx> for ConstProp {
 
         trace!("ConstProp starting for {:?}", def_id);
 
-        let dummy_body = &Body::new(
-            body.source,
-            (*body.basic_blocks).to_owned(),
-            body.source_scopes.clone(),
-            body.local_decls.clone(),
-            Default::default(),
-            body.arg_count,
-            Default::default(),
-            body.span,
-            body.generator_kind(),
-            body.tainted_by_errors,
-        );
-
         // FIXME(oli-obk, eddyb) Optimize locals (or even local paths) to hold
         // constants, instead of just checking for const-folding succeeding.
         // That would require a uniform one-def no-mutation analysis
         // and RPO (or recursing when needing the value of a local).
-        let mut optimization_finder = ConstPropagator::new(body, dummy_body, tcx);
+        let mut optimization_finder = ConstPropagator::new(body, tcx);
 
         // Traverse the body in reverse post-order, to ensure that `FullConstProp` locals are
         // assigned before being read.
-        let rpo = body.basic_blocks.reverse_postorder().to_vec();
+        for &bb in body.basic_blocks.reverse_postorder() {
-        for bb in rpo {
+            let data = &body.basic_blocks[bb];
-            let data = &mut body.basic_blocks.as_mut_preserves_cfg()[bb];
             optimization_finder.visit_basic_block_data(bb, data);
         }
 
+        let mut patch = optimization_finder.patch;
+        patch.visit_body_preserves_cfg(body);
+
         trace!("ConstProp done for {:?}", def_id);
     }
 }
@@ -143,8 +133,11 @@ impl ConstPropMachine<'_, '_> {
 
 impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for ConstPropMachine<'mir, 'tcx> {
     compile_time_machine!(<'mir, 'tcx>);
+
     const PANIC_ON_ALLOC_FAIL: bool = true; // all allocations are small (see `MAX_ALLOC_LIMIT`)
 
+    const POST_MONO_CHECKS: bool = false; // this MIR is still generic!
+
     type MemoryKind = !;
 
     #[inline(always)]
@@ -299,6 +292,7 @@ struct ConstPropagator<'mir, 'tcx> {
     tcx: TyCtxt<'tcx>,
     param_env: ParamEnv<'tcx>,
     local_decls: &'mir IndexSlice<Local, LocalDecl<'tcx>>,
+    patch: Patch<'tcx>,
 }
 
 impl<'tcx> LayoutOfHelpers<'tcx> for ConstPropagator<'_, 'tcx> {
@@ -332,11 +326,7 @@ impl<'tcx> ty::layout::HasParamEnv<'tcx> for ConstPropagator<'_, 'tcx> {
 }
 
 impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
-    fn new(
+    fn new(body: &'mir Body<'tcx>, tcx: TyCtxt<'tcx>) -> ConstPropagator<'mir, 'tcx> {
-        body: &Body<'tcx>,
-        dummy_body: &'mir Body<'tcx>,
-        tcx: TyCtxt<'tcx>,
-    ) -> ConstPropagator<'mir, 'tcx> {
         let def_id = body.source.def_id();
         let args = &GenericArgs::identity_for_item(tcx, def_id);
         let param_env = tcx.param_env_reveal_all_normalized(def_id);
@@ -367,7 +357,7 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
 
         ecx.push_stack_frame(
             Instance::new(def_id, args),
-            dummy_body,
+            body,
             &ret,
             StackPopCleanup::Root { cleanup: false },
         )
@@ -382,7 +372,8 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
             ecx.frame_mut().locals[local].make_live_uninit();
         }
 
-        ConstPropagator { ecx, tcx, param_env, local_decls: &dummy_body.local_decls }
+        let patch = Patch::new(tcx);
+        ConstPropagator { ecx, tcx, param_env, local_decls: &body.local_decls, patch }
     }
 
     fn get_const(&self, place: Place<'tcx>) -> Option<OpTy<'tcx>> {
@@ -419,12 +410,6 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
         ecx.machine.written_only_inside_own_block_locals.remove(&local);
     }
 
-    fn propagate_operand(&mut self, operand: &mut Operand<'tcx>) {
-        if let Some(place) = operand.place() && let Some(op) = self.replace_with_const(place) {
-            *operand = op;
-        }
-    }
-
     fn check_rvalue(&mut self, rvalue: &Rvalue<'tcx>) -> Option<()> {
         // Perform any special handling for specific Rvalue types.
         // Generally, checks here fall into one of two categories:
@@ -540,16 +525,7 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
         }
     }
 
-    /// Creates a new `Operand::Constant` from a `Scalar` value
+    fn replace_with_const(&mut self, place: Place<'tcx>) -> Option<ConstantKind<'tcx>> {
-    fn operand_from_scalar(&self, scalar: Scalar, ty: Ty<'tcx>) -> Operand<'tcx> {
-        Operand::Constant(Box::new(Constant {
-            span: DUMMY_SP,
-            user_ty: None,
-            literal: ConstantKind::from_scalar(self.tcx, scalar, ty),
-        }))
-    }
-
-    fn replace_with_const(&mut self, place: Place<'tcx>) -> Option<Operand<'tcx>> {
         // 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).
@@ -565,7 +541,7 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
         let Right(imm) = imm else { return None };
         match *imm {
             Immediate::Scalar(scalar) if scalar.try_to_int().is_ok() => {
-                Some(self.operand_from_scalar(scalar, value.layout.ty))
+                Some(ConstantKind::from_scalar(self.tcx, scalar, value.layout.ty))
             }
             Immediate::ScalarPair(l, r) if l.try_to_int().is_ok() && r.try_to_int().is_ok() => {
                 let alloc = self
@@ -575,15 +551,10 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
                     })
                     .ok()?;
 
-                let literal = ConstantKind::Val(
+                Some(ConstantKind::Val(
                     ConstValue::ByRef { alloc, offset: Size::ZERO },
                     value.layout.ty,
-                );
+                ))
-                Some(Operand::Constant(Box::new(Constant {
-                    span: DUMMY_SP,
-                    user_ty: None,
-                    literal,
-                })))
             }
             // Scalars or scalar pairs that contain undef values are assumed to not have
             // successfully evaluated and are thus not propagated.
@@ -725,40 +696,29 @@ impl<'tcx> Visitor<'tcx> for CanConstProp {
     }
 }
 
-impl<'tcx> MutVisitor<'tcx> for ConstPropagator<'_, 'tcx> {
+impl<'tcx> Visitor<'tcx> for ConstPropagator<'_, 'tcx> {
-    fn tcx(&self) -> TyCtxt<'tcx> {
+    fn visit_operand(&mut self, operand: &Operand<'tcx>, location: Location) {
-        self.tcx
-    }
-
-    fn visit_operand(&mut self, operand: &mut Operand<'tcx>, location: Location) {
         self.super_operand(operand, location);
-        self.propagate_operand(operand)
+        if let Some(place) = operand.place() && let Some(value) = self.replace_with_const(place) {
-    }
+            self.patch.before_effect.insert((location, place), value);
-
-    fn process_projection_elem(
-        &mut self,
-        elem: PlaceElem<'tcx>,
-        _: Location,
-    ) -> Option<PlaceElem<'tcx>> {
-        if let PlaceElem::Index(local) = elem
-            && let Some(value) = self.get_const(local.into())
-            && let Some(imm) = value.as_mplace_or_imm().right()
-            && let Immediate::Scalar(scalar) = *imm
-            && let Ok(offset) = scalar.to_target_usize(&self.tcx)
-            && let Some(min_length) = offset.checked_add(1)
-        {
-            Some(PlaceElem::ConstantIndex { offset, min_length, from_end: false })
-        } else {
-            None
         }
     }
 
-    fn visit_assign(
+    fn visit_projection_elem(
         &mut self,
-        place: &mut Place<'tcx>,
+        _: PlaceRef<'tcx>,
-        rvalue: &mut Rvalue<'tcx>,
+        elem: PlaceElem<'tcx>,
+        _: PlaceContext,
         location: Location,
     ) {
+        if let PlaceElem::Index(local) = elem
+            && let Some(value) = self.replace_with_const(local.into())
+        {
+            self.patch.before_effect.insert((location, local.into()), value);
+        }
+    }
+
+    fn visit_assign(&mut self, place: &Place<'tcx>, rvalue: &Rvalue<'tcx>, location: Location) {
         self.super_assign(place, rvalue, location);
 
         let Some(()) = self.check_rvalue(rvalue) else { return };
@@ -775,7 +735,7 @@ impl<'tcx> MutVisitor<'tcx> for ConstPropagator<'_, 'tcx> {
                     {
                         trace!("skipping replace of Rvalue::Use({:?} because it is already a const", c);
                     } else if let Some(operand) = self.replace_with_const(*place) {
-                        *rvalue = Rvalue::Use(operand);
+                        self.patch.assignments.insert(location, operand);
                     }
                 } else {
                     // Const prop failed, so erase the destination, ensuring that whatever happens
@@ -799,7 +759,7 @@ impl<'tcx> MutVisitor<'tcx> for ConstPropagator<'_, 'tcx> {
         }
     }
 
-    fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) {
+    fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) {
         trace!("visit_statement: {:?}", statement);
 
         // We want to evaluate operands before any change to the assigned-to value,
@@ -843,7 +803,7 @@ impl<'tcx> MutVisitor<'tcx> for ConstPropagator<'_, 'tcx> {
         }
     }
 
-    fn visit_basic_block_data(&mut self, block: BasicBlock, data: &mut BasicBlockData<'tcx>) {
+    fn visit_basic_block_data(&mut self, block: BasicBlock, data: &BasicBlockData<'tcx>) {
         self.super_basic_block_data(block, data);
 
         // We remove all Locals which are restricted in propagation to their containing blocks and

compiler/rustc_mir_transform/src/dataflow_const_prop.rs

@@ -7,10 +7,10 @@ use rustc_const_eval::interpret::{ImmTy, Immediate, InterpCx, OpTy, Projectable}
 use rustc_data_structures::fx::FxHashMap;
 use rustc_hir::def::DefKind;
 use rustc_middle::mir::interpret::{AllocId, ConstAllocation, ConstValue, InterpResult, Scalar};
-use rustc_middle::mir::visit::{MutVisitor, NonMutatingUseContext, PlaceContext, Visitor};
+use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor};
 use rustc_middle::mir::*;
 use rustc_middle::ty::layout::TyAndLayout;
-use rustc_middle::ty::{self, ScalarInt, Ty, TyCtxt};
+use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_mir_dataflow::value_analysis::{
     Map, PlaceIndex, State, TrackElem, ValueAnalysis, ValueAnalysisWrapper, ValueOrPlace,
 };
@@ -60,13 +60,10 @@ impl<'tcx> MirPass<'tcx> for DataflowConstProp {
             .in_scope(|| analysis.wrap().into_engine(tcx, body).iterate_to_fixpoint());
 
         // Collect results and patch the body afterwards.
-        let mut visitor = CollectAndPatch::new(tcx, &body.local_decls);
+        let mut visitor = Collector::new(tcx, &body.local_decls);
         debug_span!("collect").in_scope(|| results.visit_reachable_with(body, &mut visitor));
-        debug_span!("patch").in_scope(|| {
+        let mut patch = visitor.patch;
-            for (block, bbdata) in body.basic_blocks.as_mut_preserves_cfg().iter_enumerated_mut() {
+        debug_span!("patch").in_scope(|| patch.visit_body_preserves_cfg(body));
-                visitor.visit_basic_block_data(block, bbdata);
-            }
-        })
     }
 }
 
@@ -515,41 +512,55 @@ impl<'a, 'tcx> ConstAnalysis<'a, 'tcx> {
     }
 }
 
-struct CollectAndPatch<'tcx, 'locals> {
+pub(crate) struct Patch<'tcx> {
     tcx: TyCtxt<'tcx>,
-    local_decls: &'locals LocalDecls<'tcx>,
 
     /// For a given MIR location, this stores the values of the operands used by that location. In
     /// particular, this is before the effect, such that the operands of `_1 = _1 + _2` are
     /// properly captured. (This may become UB soon, but it is currently emitted even by safe code.)
-    before_effect: FxHashMap<(Location, Place<'tcx>), ScalarInt>,
+    pub(crate) before_effect: FxHashMap<(Location, Place<'tcx>), ConstantKind<'tcx>>,
 
     /// Stores the assigned values for assignments where the Rvalue is constant.
-    assignments: FxHashMap<Location, ScalarInt>,
+    pub(crate) assignments: FxHashMap<Location, ConstantKind<'tcx>>,
 }
 
-impl<'tcx, 'locals> CollectAndPatch<'tcx, 'locals> {
+impl<'tcx> Patch<'tcx> {
-    fn new(tcx: TyCtxt<'tcx>, local_decls: &'locals LocalDecls<'tcx>) -> Self {
+    pub(crate) fn new(tcx: TyCtxt<'tcx>) -> Self {
-        Self {
+        Self { tcx, before_effect: FxHashMap::default(), assignments: FxHashMap::default() }
-            tcx,
-            local_decls,
-            before_effect: FxHashMap::default(),
-            assignments: FxHashMap::default(),
-        }
     }
 
-    fn make_operand(&self, scalar: ScalarInt, ty: Ty<'tcx>) -> Operand<'tcx> {
+    fn make_operand(&self, literal: ConstantKind<'tcx>) -> Operand<'tcx> {
-        Operand::Constant(Box::new(Constant {
+        Operand::Constant(Box::new(Constant { span: DUMMY_SP, user_ty: None, literal }))
-            span: DUMMY_SP,
+    }
-            user_ty: None,
+}
-            literal: ConstantKind::Val(ConstValue::Scalar(scalar.into()), ty),
+
-        }))
+struct Collector<'tcx, 'locals> {
+    patch: Patch<'tcx>,
+    local_decls: &'locals LocalDecls<'tcx>,
+}
+
+impl<'tcx, 'locals> Collector<'tcx, 'locals> {
+    pub(crate) fn new(tcx: TyCtxt<'tcx>, local_decls: &'locals LocalDecls<'tcx>) -> Self {
+        Self { patch: Patch::new(tcx), local_decls }
+    }
+
+    fn try_make_constant(
+        &self,
+        place: Place<'tcx>,
+        state: &State<FlatSet<Scalar>>,
+        map: &Map,
+    ) -> Option<ConstantKind<'tcx>> {
+        let FlatSet::Elem(Scalar::Int(value)) = state.get(place.as_ref(), &map) else {
+            return None;
+        };
+        let ty = place.ty(self.local_decls, self.patch.tcx).ty;
+        Some(ConstantKind::Val(ConstValue::Scalar(value.into()), ty))
     }
 }
 
 impl<'mir, 'tcx>
     ResultsVisitor<'mir, 'tcx, Results<'tcx, ValueAnalysisWrapper<ConstAnalysis<'_, 'tcx>>>>
-    for CollectAndPatch<'tcx, '_>
+    for Collector<'tcx, '_>
 {
     type FlowState = State<FlatSet<Scalar>>;
 
@@ -581,10 +592,8 @@ impl<'mir, 'tcx>
                 // Don't overwrite the assignment if it already uses a constant (to keep the span).
             }
             StatementKind::Assign(box (place, _)) => {
-                if let FlatSet::Elem(Scalar::Int(value)) =
+                if let Some(value) = self.try_make_constant(place, state, &results.analysis.0.map) {
-                    state.get(place.as_ref(), &results.analysis.0.map)
+                    self.patch.assignments.insert(location, value);
-                {
-                    self.assignments.insert(location, value);
                 }
             }
             _ => (),
@@ -603,7 +612,7 @@ impl<'mir, 'tcx>
     }
 }
 
-impl<'tcx> MutVisitor<'tcx> for CollectAndPatch<'tcx, '_> {
+impl<'tcx> MutVisitor<'tcx> for Patch<'tcx> {
     fn tcx(&self) -> TyCtxt<'tcx> {
         self.tcx
     }
@@ -612,8 +621,7 @@ impl<'tcx> MutVisitor<'tcx> for CollectAndPatch<'tcx, '_> {
         if let Some(value) = self.assignments.get(&location) {
             match &mut statement.kind {
                 StatementKind::Assign(box (_, rvalue)) => {
-                    let ty = rvalue.ty(self.local_decls, self.tcx);
+                    *rvalue = Rvalue::Use(self.make_operand(*value));
-                    *rvalue = Rvalue::Use(self.make_operand(*value, ty));
                 }
                 _ => bug!("found assignment info for non-assign statement"),
             }
@@ -626,8 +634,7 @@ impl<'tcx> MutVisitor<'tcx> for CollectAndPatch<'tcx, '_> {
         match operand {
             Operand::Copy(place) | Operand::Move(place) => {
                 if let Some(value) = self.before_effect.get(&(location, *place)) {
-                    let ty = place.ty(self.local_decls, self.tcx).ty;
+                    *operand = self.make_operand(*value);
-                    *operand = self.make_operand(*value, ty);
                 } else if !place.projection.is_empty() {
                     self.super_operand(operand, location)
                 }
@@ -641,11 +648,11 @@ impl<'tcx> MutVisitor<'tcx> for CollectAndPatch<'tcx, '_> {
         elem: PlaceElem<'tcx>,
         location: Location,
     ) -> Option<PlaceElem<'tcx>> {
-        if let PlaceElem::Index(local) = elem
+        if let PlaceElem::Index(local) = elem {
-            && let Some(value) = self.before_effect.get(&(location, local.into()))
+            let offset = self.before_effect.get(&(location, local.into()))?;
-            && let Ok(offset) = value.try_to_target_usize(self.tcx)
+            let offset = offset.try_to_scalar()?;
-            && let Some(min_length) = offset.checked_add(1)
+            let offset = offset.to_target_usize(&self.tcx).ok()?;
-        {
+            let min_length = offset.checked_add(1)?;
             Some(PlaceElem::ConstantIndex { offset, min_length, from_end: false })
         } else {
             None
@@ -655,29 +662,35 @@ impl<'tcx> MutVisitor<'tcx> for CollectAndPatch<'tcx, '_> {
 
 struct OperandCollector<'tcx, 'map, 'locals, 'a> {
     state: &'a State<FlatSet<Scalar>>,
-    visitor: &'a mut CollectAndPatch<'tcx, 'locals>,
+    visitor: &'a mut Collector<'tcx, 'locals>,
     map: &'map Map,
 }
 
 impl<'tcx> Visitor<'tcx> for OperandCollector<'tcx, '_, '_, '_> {
+    fn visit_projection_elem(
+        &mut self,
+        _: PlaceRef<'tcx>,
+        elem: PlaceElem<'tcx>,
+        _: PlaceContext,
+        location: Location,
+    ) {
+        if let PlaceElem::Index(local) = elem
+            && let Some(value) = self.visitor.try_make_constant(local.into(), self.state, self.map)
+        {
+            self.visitor.patch.before_effect.insert((location, local.into()), value);
+        }
+    }
+
     fn visit_operand(&mut self, operand: &Operand<'tcx>, location: Location) {
         if let Some(place) = operand.place() {
-            if let FlatSet::Elem(Scalar::Int(value)) = self.state.get(place.as_ref(), self.map) {
+            if let Some(value) = self.visitor.try_make_constant(place, self.state, self.map) {
-                self.visitor.before_effect.insert((location, place), value);
+                self.visitor.patch.before_effect.insert((location, place), value);
             } else if !place.projection.is_empty() {
                 // Try to propagate into `Index` projections.
                 self.super_operand(operand, location)
             }
         }
     }
-
-    fn visit_local(&mut self, local: Local, ctxt: PlaceContext, location: Location) {
-        if let PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy | NonMutatingUseContext::Move) = ctxt
-            && let FlatSet::Elem(Scalar::Int(value)) = self.state.get(local.into(), self.map)
-        {
-            self.visitor.before_effect.insert((location, local.into()), value);
-        }
-    }
 }
 
 struct DummyMachine;