From 32fe00ae85adfd8dd7c603be93ed7d6c37e11dd4 Mon Sep 17 00:00:00 2001 From: Camille GILLOT Date: Thu, 3 Aug 2023 12:06:37 +0000 Subject: [PATCH] Reuse CollectAndPatch for normal ConstProp. --- .../rustc_mir_transform/src/const_prop.rs | 86 +++++++------------ .../src/dataflow_const_prop.rs | 82 ++++++++++-------- 2 files changed, 75 insertions(+), 93 deletions(-) diff --git a/compiler/rustc_mir_transform/src/const_prop.rs b/compiler/rustc_mir_transform/src/const_prop.rs index 75e5c8a8d40..5c54080a173 100644 --- 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>, + 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> { @@ -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> { + fn replace_with_const(&mut self, place: Place<'tcx>) -> Option> { // 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) - } - - fn process_projection_elem( - &mut self, - elem: PlaceElem<'tcx>, - _: Location, - ) -> Option> { - 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 + if let Some(place) = operand.place() && let Some(value) = self.replace_with_const(place) { + self.patch.before_effect.insert((location, place), value); } } - fn visit_assign( + fn visit_projection_elem( &mut self, - place: &mut Place<'tcx>, - rvalue: &mut Rvalue<'tcx>, + _: PlaceRef<'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 }; @@ -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 index 5b4fbe54e4a..bff6044ecca 100644 --- 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>, + pub(crate) assignments: FxHashMap>, } -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>>> - for CollectAndPatch<'tcx, '_> + for Collector<'tcx, '_> { type FlowState = State>; @@ -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>, - 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;