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