compiler/rustc_mir_transform/src/dataflow_const_prop.rs

   1 //! A constant propagation optimization pass based on dataflow analysis.
   2 //!
   3 //! Currently, this pass only propagates scalar values.
   4
   5 use rustc_const_eval::interpret::{ConstValue, ImmTy, Immediate, InterpCx, Scalar};
   6 use rustc_data_structures::fx::FxHashMap;
   7 use rustc_middle::mir::visit::{MutVisitor, Visitor};
   8 use rustc_middle::mir::*;
   9 use rustc_middle::ty::{self, Ty, TyCtxt};
  10 use rustc_mir_dataflow::value_analysis::{Map, State, TrackElem, ValueAnalysis, ValueOrPlace};
  11 use rustc_mir_dataflow::{lattice::FlatSet, Analysis, ResultsVisitor, SwitchIntEdgeEffects};
  12 use rustc_span::DUMMY_SP;
  13
  14 use crate::MirPass;
  15
  16 // These constants are somewhat random guesses and have not been optimized.
  17 // If `tcx.sess.mir_opt_level() >= 4`, we ignore the limits (this can become very expensive).
  18 const BLOCK_LIMIT: usize = 100;
  19 const PLACE_LIMIT: usize = 100;
  20
  21 pub struct DataflowConstProp;
  22
  23 impl<'tcx> MirPass<'tcx> for DataflowConstProp {
  24     fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
  25         sess.mir_opt_level() >= 3
  26     }
  27
  28     #[instrument(skip_all level = "debug")]
  29     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
  30         if tcx.sess.mir_opt_level() < 4 && body.basic_blocks.len() > BLOCK_LIMIT {
  31             debug!("aborted dataflow const prop due too many basic blocks");
  32             return;
  33         }
  34
  35         // Decide which places to track during the analysis.
  36         let map = Map::from_filter(tcx, body, Ty::is_scalar);
  37
  38         // We want to have a somewhat linear runtime w.r.t. the number of statements/terminators.
  39         // Let's call this number `n`. Dataflow analysis has `O(h*n)` transfer function
  40         // applications, where `h` is the height of the lattice. Because the height of our lattice
  41         // is linear w.r.t. the number of tracked places, this is `O(tracked_places * n)`. However,
  42         // because every transfer function application could traverse the whole map, this becomes
  43         // `O(num_nodes * tracked_places * n)` in terms of time complexity. Since the number of
  44         // map nodes is strongly correlated to the number of tracked places, this becomes more or
  45         // less `O(n)` if we place a constant limit on the number of tracked places.
  46         if tcx.sess.mir_opt_level() < 4 && map.tracked_places() > PLACE_LIMIT {
  47             debug!("aborted dataflow const prop due to too many tracked places");
  48             return;
  49         }
  50
  51         // Perform the actual dataflow analysis.
  52         let analysis = ConstAnalysis::new(tcx, body, map);
  53         let results = debug_span!("analyze")
  54             .in_scope(|| analysis.wrap().into_engine(tcx, body).iterate_to_fixpoint());
  55
  56         // Collect results and patch the body afterwards.
  57         let mut visitor = CollectAndPatch::new(tcx, &results.analysis.0.map);
  58         debug_span!("collect").in_scope(|| results.visit_reachable_with(body, &mut visitor));
  59         debug_span!("patch").in_scope(|| visitor.visit_body(body));
  60     }
  61 }
  62
  63 struct ConstAnalysis<'tcx> {
  64     map: Map,
  65     tcx: TyCtxt<'tcx>,
  66     ecx: InterpCx<'tcx, 'tcx, DummyMachine>,
  67     param_env: ty::ParamEnv<'tcx>,
  68 }
  69
  70 impl<'tcx> ValueAnalysis<'tcx> for ConstAnalysis<'tcx> {
  71     type Value = FlatSet<ScalarTy<'tcx>>;
  72
  73     const NAME: &'static str = "ConstAnalysis";
  74
  75     fn map(&self) -> &Map {
  76         &self.map
  77     }
  78
  79     fn handle_assign(
  80         &self,
  81         target: Place<'tcx>,
  82         rvalue: &Rvalue<'tcx>,
  83         state: &mut State<Self::Value>,
  84     ) {
  85         match rvalue {
  86             Rvalue::CheckedBinaryOp(op, box (left, right)) => {
  87                 let target = self.map().find(target.as_ref());
  88                 if let Some(target) = target {
  89                     // We should not track any projections other than
  90                     // what is overwritten below, but just in case...
  91                     state.flood_idx(target, self.map());
  92                 }
  93
  94                 let value_target = target
  95                     .and_then(|target| self.map().apply(target, TrackElem::Field(0_u32.into())));
  96                 let overflow_target = target
  97                     .and_then(|target| self.map().apply(target, TrackElem::Field(1_u32.into())));
  98
  99                 if value_target.is_some() || overflow_target.is_some() {
 100                     let (val, overflow) = self.binary_op(state, *op, left, right);
 101
 102                     if let Some(value_target) = value_target {
 103                         state.assign_idx(value_target, ValueOrPlace::Value(val), self.map());
 104                     }
 105                     if let Some(overflow_target) = overflow_target {
 106                         let overflow = match overflow {
 107                             FlatSet::Top => FlatSet::Top,
 108                             FlatSet::Elem(overflow) => {
 109                                 if overflow {
 110                                     // Overflow cannot be reliably propagated. See: https://github.com/rust-lang/rust/pull/101168#issuecomment-1288091446
 111                                     FlatSet::Top
 112                                 } else {
 113                                     self.wrap_scalar(Scalar::from_bool(false), self.tcx.types.bool)
 114                                 }
 115                             }
 116                             FlatSet::Bottom => FlatSet::Bottom,
 117                         };
 118                         state.assign_idx(
 119                             overflow_target,
 120                             ValueOrPlace::Value(overflow),
 121                             self.map(),
 122                         );
 123                     }
 124                 }
 125             }
 126             _ => self.super_assign(target, rvalue, state),
 127         }
 128     }
 129
 130     fn handle_rvalue(
 131         &self,
 132         rvalue: &Rvalue<'tcx>,
 133         state: &mut State<Self::Value>,
 134     ) -> ValueOrPlace<Self::Value> {
 135         match rvalue {
 136             Rvalue::Cast(
 137                 kind @ (CastKind::IntToInt
 138                 | CastKind::FloatToInt
 139                 | CastKind::FloatToFloat
 140                 | CastKind::IntToFloat),
 141                 operand,
 142                 ty,
 143             ) => match self.eval_operand(operand, state) {
 144                 FlatSet::Elem(op) => match kind {
 145                     CastKind::IntToInt | CastKind::IntToFloat => {
 146                         self.ecx.int_to_int_or_float(&op, *ty)
 147                     }
 148                     CastKind::FloatToInt | CastKind::FloatToFloat => {
 149                         self.ecx.float_to_float_or_int(&op, *ty)
 150                     }
 151                     _ => unreachable!(),
 152                 }
 153                 .map(|result| ValueOrPlace::Value(self.wrap_immediate(result, *ty)))
 154                 .unwrap_or(ValueOrPlace::top()),
 155                 _ => ValueOrPlace::top(),
 156             },
 157             Rvalue::BinaryOp(op, box (left, right)) => {
 158                 // Overflows must be ignored here.
 159                 let (val, _overflow) = self.binary_op(state, *op, left, right);
 160                 ValueOrPlace::Value(val)
 161             }
 162             Rvalue::UnaryOp(op, operand) => match self.eval_operand(operand, state) {
 163                 FlatSet::Elem(value) => self
 164                     .ecx
 165                     .unary_op(*op, &value)
 166                     .map(|val| ValueOrPlace::Value(self.wrap_immty(val)))
 167                     .unwrap_or(ValueOrPlace::Value(FlatSet::Top)),
 168                 FlatSet::Bottom => ValueOrPlace::Value(FlatSet::Bottom),
 169                 FlatSet::Top => ValueOrPlace::Value(FlatSet::Top),
 170             },
 171             _ => self.super_rvalue(rvalue, state),
 172         }
 173     }
 174
 175     fn handle_constant(
 176         &self,
 177         constant: &Constant<'tcx>,
 178         _state: &mut State<Self::Value>,
 179     ) -> Self::Value {
 180         constant
 181             .literal
 182             .eval(self.tcx, self.param_env)
 183             .try_to_scalar()
 184             .map(|value| FlatSet::Elem(ScalarTy(value, constant.ty())))
 185             .unwrap_or(FlatSet::Top)
 186     }
 187
 188     fn handle_switch_int(
 189         &self,
 190         discr: &Operand<'tcx>,
 191         apply_edge_effects: &mut impl SwitchIntEdgeEffects<State<Self::Value>>,
 192     ) {
 193         // FIXME: The dataflow framework only provides the state if we call `apply()`, which makes
 194         // this more inefficient than it has to be.
 195         let mut discr_value = None;
 196         let mut handled = false;
 197         apply_edge_effects.apply(|state, target| {
 198             let discr_value = match discr_value {
 199                 Some(value) => value,
 200                 None => {
 201                     let value = match self.handle_operand(discr, state) {
 202                         ValueOrPlace::Value(value) => value,
 203                         ValueOrPlace::Place(place) => state.get_idx(place, self.map()),
 204                     };
 205                     let result = match value {
 206                         FlatSet::Top => FlatSet::Top,
 207                         FlatSet::Elem(ScalarTy(scalar, _)) => {
 208                             let int = scalar.assert_int();
 209                             FlatSet::Elem(int.assert_bits(int.size()))
 210                         }
 211                         FlatSet::Bottom => FlatSet::Bottom,
 212                     };
 213                     discr_value = Some(result);
 214                     result
 215                 }
 216             };
 217
 218             let FlatSet::Elem(choice) = discr_value else {
 219                 // Do nothing if we don't know which branch will be taken.
 220                 return
 221             };
 222
 223             if target.value.map(|n| n == choice).unwrap_or(!handled) {
 224                 // Branch is taken. Has no effect on state.
 225                 handled = true;
 226             } else {
 227                 // Branch is not taken.
 228                 state.mark_unreachable();
 229             }
 230         })
 231     }
 232 }
 233
 234 #[derive(Clone, PartialEq, Eq)]
 235 struct ScalarTy<'tcx>(Scalar, Ty<'tcx>);
 236
 237 impl<'tcx> std::fmt::Debug for ScalarTy<'tcx> {
 238     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
 239         // This is used for dataflow visualization, so we return something more concise.
 240         std::fmt::Display::fmt(&ConstantKind::Val(ConstValue::Scalar(self.0), self.1), f)
 241     }
 242 }
 243
 244 impl<'tcx> ConstAnalysis<'tcx> {
 245     pub fn new(tcx: TyCtxt<'tcx>, body: &Body<'tcx>, map: Map) -> Self {
 246         let param_env = tcx.param_env(body.source.def_id());
 247         Self {
 248             map,
 249             tcx,
 250             ecx: InterpCx::new(tcx, DUMMY_SP, param_env, DummyMachine),
 251             param_env: param_env,
 252         }
 253     }
 254
 255     fn binary_op(
 256         &self,
 257         state: &mut State<FlatSet<ScalarTy<'tcx>>>,
 258         op: BinOp,
 259         left: &Operand<'tcx>,
 260         right: &Operand<'tcx>,
 261     ) -> (FlatSet<ScalarTy<'tcx>>, FlatSet<bool>) {
 262         let left = self.eval_operand(left, state);
 263         let right = self.eval_operand(right, state);
 264         match (left, right) {
 265             (FlatSet::Elem(left), FlatSet::Elem(right)) => {
 266                 match self.ecx.overflowing_binary_op(op, &left, &right) {
 267                     Ok((val, overflow, ty)) => (self.wrap_scalar(val, ty), FlatSet::Elem(overflow)),
 268                     _ => (FlatSet::Top, FlatSet::Top),
 269                 }
 270             }
 271             (FlatSet::Bottom, _) | (_, FlatSet::Bottom) => (FlatSet::Bottom, FlatSet::Bottom),
 272             (_, _) => {
 273                 // Could attempt some algebraic simplifcations here.
 274                 (FlatSet::Top, FlatSet::Top)
 275             }
 276         }
 277     }
 278
 279     fn eval_operand(
 280         &self,
 281         op: &Operand<'tcx>,
 282         state: &mut State<FlatSet<ScalarTy<'tcx>>>,
 283     ) -> FlatSet<ImmTy<'tcx>> {
 284         let value = match self.handle_operand(op, state) {
 285             ValueOrPlace::Value(value) => value,
 286             ValueOrPlace::Place(place) => state.get_idx(place, &self.map),
 287         };
 288         match value {
 289             FlatSet::Top => FlatSet::Top,
 290             FlatSet::Elem(ScalarTy(scalar, ty)) => self
 291                 .tcx
 292                 .layout_of(self.param_env.and(ty))
 293                 .map(|layout| FlatSet::Elem(ImmTy::from_scalar(scalar, layout)))
 294                 .unwrap_or(FlatSet::Top),
 295             FlatSet::Bottom => FlatSet::Bottom,
 296         }
 297     }
 298
 299     fn wrap_scalar(&self, scalar: Scalar, ty: Ty<'tcx>) -> FlatSet<ScalarTy<'tcx>> {
 300         FlatSet::Elem(ScalarTy(scalar, ty))
 301     }
 302
 303     fn wrap_immediate(&self, imm: Immediate, ty: Ty<'tcx>) -> FlatSet<ScalarTy<'tcx>> {
 304         match imm {
 305             Immediate::Scalar(scalar) => self.wrap_scalar(scalar, ty),
 306             _ => FlatSet::Top,
 307         }
 308     }
 309
 310     fn wrap_immty(&self, val: ImmTy<'tcx>) -> FlatSet<ScalarTy<'tcx>> {
 311         self.wrap_immediate(*val, val.layout.ty)
 312     }
 313 }
 314
 315 struct CollectAndPatch<'tcx, 'map> {
 316     tcx: TyCtxt<'tcx>,
 317     map: &'map Map,
 318
 319     /// For a given MIR location, this stores the values of the operands used by that location. In
 320     /// particular, this is before the effect, such that the operands of `_1 = _1 + _2` are
 321     /// properly captured. (This may become UB soon, but it is currently emitted even by safe code.)
 322     before_effect: FxHashMap<(Location, Place<'tcx>), ScalarTy<'tcx>>,
 323
 324     /// Stores the assigned values for assignments where the Rvalue is constant.
 325     assignments: FxHashMap<Location, ScalarTy<'tcx>>,
 326 }
 327
 328 impl<'tcx, 'map> CollectAndPatch<'tcx, 'map> {
 329     fn new(tcx: TyCtxt<'tcx>, map: &'map Map) -> Self {
 330         Self { tcx, map, before_effect: FxHashMap::default(), assignments: FxHashMap::default() }
 331     }
 332
 333     fn make_operand(&self, scalar: ScalarTy<'tcx>) -> Operand<'tcx> {
 334         Operand::Constant(Box::new(Constant {
 335             span: DUMMY_SP,
 336             user_ty: None,
 337             literal: ConstantKind::Val(ConstValue::Scalar(scalar.0), scalar.1),
 338         }))
 339     }
 340 }
 341
 342 impl<'mir, 'tcx, 'map> ResultsVisitor<'mir, 'tcx> for CollectAndPatch<'tcx, 'map> {
 343     type FlowState = State<FlatSet<ScalarTy<'tcx>>>;
 344
 345     fn visit_statement_before_primary_effect(
 346         &mut self,
 347         state: &Self::FlowState,
 348         statement: &'mir Statement<'tcx>,
 349         location: Location,
 350     ) {
 351         match &statement.kind {
 352             StatementKind::Assign(box (_, rvalue)) => {
 353                 OperandCollector { state, visitor: self }.visit_rvalue(rvalue, location);
 354             }
 355             _ => (),
 356         }
 357     }
 358
 359     fn visit_statement_after_primary_effect(
 360         &mut self,
 361         state: &Self::FlowState,
 362         statement: &'mir Statement<'tcx>,
 363         location: Location,
 364     ) {
 365         match statement.kind {
 366             StatementKind::Assign(box (_, Rvalue::Use(Operand::Constant(_)))) => {
 367                 // Don't overwrite the assignment if it already uses a constant (to keep the span).
 368             }
 369             StatementKind::Assign(box (place, _)) => match state.get(place.as_ref(), self.map) {
 370                 FlatSet::Top => (),
 371                 FlatSet::Elem(value) => {
 372                     self.assignments.insert(location, value);
 373                 }
 374                 FlatSet::Bottom => {
 375                     // This assignment is either unreachable, or an uninitialized value is assigned.
 376                 }
 377             },
 378             _ => (),
 379         }
 380     }
 381
 382     fn visit_terminator_before_primary_effect(
 383         &mut self,
 384         state: &Self::FlowState,
 385         terminator: &'mir Terminator<'tcx>,
 386         location: Location,
 387     ) {
 388         OperandCollector { state, visitor: self }.visit_terminator(terminator, location);
 389     }
 390 }
 391
 392 impl<'tcx, 'map> MutVisitor<'tcx> for CollectAndPatch<'tcx, 'map> {
 393     fn tcx<'a>(&'a self) -> TyCtxt<'tcx> {
 394         self.tcx
 395     }
 396
 397     fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) {
 398         if let Some(value) = self.assignments.get(&location) {
 399             match &mut statement.kind {
 400                 StatementKind::Assign(box (_, rvalue)) => {
 401                     *rvalue = Rvalue::Use(self.make_operand(value.clone()));
 402                 }
 403                 _ => bug!("found assignment info for non-assign statement"),
 404             }
 405         } else {
 406             self.super_statement(statement, location);
 407         }
 408     }
 409
 410     fn visit_operand(&mut self, operand: &mut Operand<'tcx>, location: Location) {
 411         match operand {
 412             Operand::Copy(place) | Operand::Move(place) => {
 413                 if let Some(value) = self.before_effect.get(&(location, *place)) {
 414                     *operand = self.make_operand(value.clone());
 415                 }
 416             }
 417             _ => (),
 418         }
 419     }
 420 }
 421
 422 struct OperandCollector<'tcx, 'map, 'a> {
 423     state: &'a State<FlatSet<ScalarTy<'tcx>>>,
 424     visitor: &'a mut CollectAndPatch<'tcx, 'map>,
 425 }
 426
 427 impl<'tcx, 'map, 'a> Visitor<'tcx> for OperandCollector<'tcx, 'map, 'a> {
 428     fn visit_operand(&mut self, operand: &Operand<'tcx>, location: Location) {
 429         match operand {
 430             Operand::Copy(place) | Operand::Move(place) => {
 431                 match self.state.get(place.as_ref(), self.visitor.map) {
 432                     FlatSet::Top => (),
 433                     FlatSet::Elem(value) => {
 434                         self.visitor.before_effect.insert((location, *place), value);
 435                     }
 436                     FlatSet::Bottom => (),
 437                 }
 438             }
 439             _ => (),
 440         }
 441     }
 442 }
 443
 444 struct DummyMachine;
 445
 446 impl<'mir, 'tcx> rustc_const_eval::interpret::Machine<'mir, 'tcx> for DummyMachine {
 447     rustc_const_eval::interpret::compile_time_machine!(<'mir, 'tcx>);
 448     type MemoryKind = !;
 449     const PANIC_ON_ALLOC_FAIL: bool = true;
 450
 451     fn enforce_alignment(_ecx: &InterpCx<'mir, 'tcx, Self>) -> bool {
 452         unimplemented!()
 453     }
 454
 455     fn enforce_validity(_ecx: &InterpCx<'mir, 'tcx, Self>) -> bool {
 456         unimplemented!()
 457     }
 458
 459     fn find_mir_or_eval_fn(
 460         _ecx: &mut InterpCx<'mir, 'tcx, Self>,
 461         _instance: ty::Instance<'tcx>,
 462         _abi: rustc_target::spec::abi::Abi,
 463         _args: &[rustc_const_eval::interpret::OpTy<'tcx, Self::Provenance>],
 464         _destination: &rustc_const_eval::interpret::PlaceTy<'tcx, Self::Provenance>,
 465         _target: Option<BasicBlock>,
 466         _unwind: rustc_const_eval::interpret::StackPopUnwind,
 467     ) -> interpret::InterpResult<'tcx, Option<(&'mir Body<'tcx>, ty::Instance<'tcx>)>> {
 468         unimplemented!()
 469     }
 470
 471     fn call_intrinsic(
 472         _ecx: &mut InterpCx<'mir, 'tcx, Self>,
 473         _instance: ty::Instance<'tcx>,
 474         _args: &[rustc_const_eval::interpret::OpTy<'tcx, Self::Provenance>],
 475         _destination: &rustc_const_eval::interpret::PlaceTy<'tcx, Self::Provenance>,
 476         _target: Option<BasicBlock>,
 477         _unwind: rustc_const_eval::interpret::StackPopUnwind,
 478     ) -> interpret::InterpResult<'tcx> {
 479         unimplemented!()
 480     }
 481
 482     fn assert_panic(
 483         _ecx: &mut InterpCx<'mir, 'tcx, Self>,
 484         _msg: &rustc_middle::mir::AssertMessage<'tcx>,
 485         _unwind: Option<BasicBlock>,
 486     ) -> interpret::InterpResult<'tcx> {
 487         unimplemented!()
 488     }
 489
 490     fn binary_ptr_op(
 491         _ecx: &InterpCx<'mir, 'tcx, Self>,
 492         _bin_op: BinOp,
 493         _left: &rustc_const_eval::interpret::ImmTy<'tcx, Self::Provenance>,
 494         _right: &rustc_const_eval::interpret::ImmTy<'tcx, Self::Provenance>,
 495     ) -> interpret::InterpResult<'tcx, (interpret::Scalar<Self::Provenance>, bool, Ty<'tcx>)> {
 496         throw_unsup!(Unsupported("".into()))
 497     }
 498
 499     fn expose_ptr(
 500         _ecx: &mut InterpCx<'mir, 'tcx, Self>,
 501         _ptr: interpret::Pointer<Self::Provenance>,
 502     ) -> interpret::InterpResult<'tcx> {
 503         unimplemented!()
 504     }
 505
 506     fn init_frame_extra(
 507         _ecx: &mut InterpCx<'mir, 'tcx, Self>,
 508         _frame: rustc_const_eval::interpret::Frame<'mir, 'tcx, Self::Provenance>,
 509     ) -> interpret::InterpResult<
 510         'tcx,
 511         rustc_const_eval::interpret::Frame<'mir, 'tcx, Self::Provenance, Self::FrameExtra>,
 512     > {
 513         unimplemented!()
 514     }
 515
 516     fn stack<'a>(
 517         _ecx: &'a InterpCx<'mir, 'tcx, Self>,
 518     ) -> &'a [rustc_const_eval::interpret::Frame<'mir, 'tcx, Self::Provenance, Self::FrameExtra>]
 519     {
 520         unimplemented!()
 521     }
 522
 523     fn stack_mut<'a>(
 524         _ecx: &'a mut InterpCx<'mir, 'tcx, Self>,
 525     ) -> &'a mut Vec<
 526         rustc_const_eval::interpret::Frame<'mir, 'tcx, Self::Provenance, Self::FrameExtra>,
 527     > {
 528         unimplemented!()
 529     }
 530 }