Reuse CollectAndPatch for normal ConstProp.

cjgillot · cjgillot · commit 32fe00ae85ad · 2023-09-12T20:17:55.000Z
diff --git a/compiler/rustc_mir_transform/src/const_prop.rs b/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,
@@ -120,6 +121,8 @@ impl<'tcx> MirPass<'tcx> for ConstProp {
             optimization_finder.visit_basic_block_data(bb, data);
         }
 
+        optimization_finder.patch.visit_body_preserves_cfg(body);
+
         trace!("ConstProp done for {:?}", def_id);
     }
 }
@@ -302,6 +305,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> {
@@ -385,7 +389,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: &dummy_body.local_decls, patch }
     }
 
     fn get_const(&self, place: Place<'tcx>) -> Option<OpTy<'tcx>> {
@@ -422,12 +427,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:
@@ -543,16 +542,7 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
         }
     }
 
-    /// Creates a new `Operand::Constant` from a `Scalar` value
-    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>> {
+    fn replace_with_const(&mut self, place: Place<'tcx>) -> Option<ConstantKind<'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).
@@ -568,7 +558,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
@@ -578,15 +568,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.
@@ -728,40 +713,29 @@ impl<'tcx> Visitor<'tcx> for CanConstProp {
     }
 }
 
-impl<'tcx> MutVisitor<'tcx> for ConstPropagator<'_, 'tcx> {
-    fn tcx(&self) -> TyCtxt<'tcx> {
-        self.tcx
-    }
-
-    fn visit_operand(&mut self, operand: &mut Operand<'tcx>, location: Location) {
+impl<'tcx> Visitor<'tcx> for ConstPropagator<'_, 'tcx> {
+    fn visit_operand(&mut self, operand: &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(
+    fn visit_projection_elem(
         &mut self,
+        _: PlaceRef<'tcx>,
         elem: PlaceElem<'tcx>,
-        _: Location,
-    ) -> Option<PlaceElem<'tcx>> {
+        _: PlaceContext,
+        location: Location,
+    ) {
         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)
+            && let Some(value) = self.replace_with_const(local.into())
         {
-            Some(PlaceElem::ConstantIndex { offset, min_length, from_end: false })
-        } else {
-            None
+            self.patch.before_effect.insert((location, local.into()), value);
         }
     }
 
-    fn visit_assign(
-        &mut self,
-        place: &mut Place<'tcx>,
-        rvalue: &mut Rvalue<'tcx>,
-        location: Location,
-    ) {
+    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 };
@@ -778,7 +752,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
@@ -802,7 +776,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,
@@ -846,7 +820,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
diff --git a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
@@ -7,7 +7,7 @@ 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, Ty, TyCtxt};
@@ -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(|| {
-            for (block, bbdata) in body.basic_blocks.as_mut_preserves_cfg().iter_enumerated_mut() {
-                visitor.visit_basic_block_data(block, bbdata);
-            }
-        })
+        let mut patch = visitor.patch;
+        debug_span!("patch").in_scope(|| patch.visit_body_preserves_cfg(body));
     }
 }
 
@@ -517,27 +514,36 @@ 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>), ConstantKind<'tcx>>,
+    pub(crate) before_effect: FxHashMap<(Location, Place<'tcx>), ConstantKind<'tcx>>,
 
     /// Stores the assigned values for assignments where the Rvalue is constant.
-    assignments: FxHashMap<Location, ConstantKind<'tcx>>,
+    pub(crate) assignments: FxHashMap<Location, ConstantKind<'tcx>>,
 }
 
-impl<'tcx, 'locals> CollectAndPatch<'tcx, 'locals> {
-    fn new(tcx: TyCtxt<'tcx>, local_decls: &'locals LocalDecls<'tcx>) -> Self {
-        Self {
-            tcx,
-            local_decls,
-            before_effect: FxHashMap::default(),
-            assignments: FxHashMap::default(),
-        }
+impl<'tcx> Patch<'tcx> {
+    pub(crate) fn new(tcx: TyCtxt<'tcx>) -> Self {
+        Self { tcx, before_effect: FxHashMap::default(), assignments: FxHashMap::default() }
+    }
+
+    fn make_operand(&self, literal: ConstantKind<'tcx>) -> Operand<'tcx> {
+        Operand::Constant(Box::new(Constant { span: DUMMY_SP, user_ty: None, literal }))
+    }
+}
+
+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(
@@ -549,18 +555,14 @@ impl<'tcx, 'locals> CollectAndPatch<'tcx, 'locals> {
         let FlatSet::Elem(Scalar::Int(value)) = state.get(place.as_ref(), &map) else {
             return None;
         };
-        let ty = place.ty(self.local_decls, self.tcx).ty;
+        let ty = place.ty(self.local_decls, self.patch.tcx).ty;
         Some(ConstantKind::Val(ConstValue::Scalar(value.into()), ty))
     }
-
-    fn make_operand(&self, literal: ConstantKind<'tcx>) -> Operand<'tcx> {
-        Operand::Constant(Box::new(Constant { span: DUMMY_SP, user_ty: None, literal }))
-    }
 }
 
 impl<'mir, 'tcx>
     ResultsVisitor<'mir, 'tcx, Results<'tcx, ValueAnalysisWrapper<ConstAnalysis<'_, 'tcx>>>>
-    for CollectAndPatch<'tcx, '_>
+    for Collector<'tcx, '_>
 {
     type FlowState = State<FlatSet<Scalar>>;
 
@@ -593,7 +595,7 @@ impl<'mir, 'tcx>
             }
             StatementKind::Assign(box (place, _)) => {
                 if let Some(value) = self.try_make_constant(place, state, &results.analysis.0.map) {
-                    self.assignments.insert(location, value);
+                    self.patch.assignments.insert(location, value);
                 }
             }
             _ => (),
@@ -612,7 +614,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
     }
@@ -662,29 +664,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 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 Some(value) = self.visitor.try_make_constant(local.into(), self.state, self.map)
-        {
-            self.visitor.before_effect.insert((location, local.into()), value);
-        }
-    }
 }
 
 struct DummyMachine;
