compiler/rustc_mir_transform/src/simplify_comparison_integral.rs

   1 use std::iter;
   2
   3 use super::MirPass;
   4 use rustc_middle::{
   5     mir::{
   6         interpret::Scalar, BasicBlock, BinOp, Body, Operand, Place, Rvalue, Statement,
   7         StatementKind, SwitchTargets, TerminatorKind,
   8     },
   9     ty::{Ty, TyCtxt},
  10 };
  11
  12 /// Pass to convert `if` conditions on integrals into switches on the integral.
  13 /// For an example, it turns something like
  14 ///
  15 /// ```ignore (MIR)
  16 /// _3 = Eq(move _4, const 43i32);
  17 /// StorageDead(_4);
  18 /// switchInt(_3) -> [false: bb2, otherwise: bb3];
  19 /// ```
  20 ///
  21 /// into:
  22 ///
  23 /// ```ignore (MIR)
  24 /// switchInt(_4) -> [43i32: bb3, otherwise: bb2];
  25 /// ```
  26 pub struct SimplifyComparisonIntegral;
  27
  28 impl<'tcx> MirPass<'tcx> for SimplifyComparisonIntegral {
  29     fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
  30         sess.mir_opt_level() > 0
  31     }
  32
  33     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
  34         trace!("Running SimplifyComparisonIntegral on {:?}", body.source);
  35
  36         let helper = OptimizationFinder { body };
  37         let opts = helper.find_optimizations();
  38         let mut storage_deads_to_insert = vec![];
  39         let mut storage_deads_to_remove: Vec<(usize, BasicBlock)> = vec![];
  40         let param_env = tcx.param_env(body.source.def_id());
  41         for opt in opts {
  42             trace!("SUCCESS: Applying {:?}", opt);
  43             // replace terminator with a switchInt that switches on the integer directly
  44             let bbs = &mut body.basic_blocks_mut();
  45             let bb = &mut bbs[opt.bb_idx];
  46             let new_value = match opt.branch_value_scalar {
  47                 Scalar::Int(int) => {
  48                     let layout = tcx
  49                         .layout_of(param_env.and(opt.branch_value_ty))
  50                         .expect("if we have an evaluated constant we must know the layout");
  51                     int.assert_bits(layout.size)
  52                 }
  53                 Scalar::Ptr(..) => continue,
  54             };
  55             const FALSE: u128 = 0;
  56
  57             let mut new_targets = opt.targets;
  58             let first_value = new_targets.iter().next().unwrap().0;
  59             let first_is_false_target = first_value == FALSE;
  60             match opt.op {
  61                 BinOp::Eq => {
  62                     // if the assignment was Eq we want the true case to be first
  63                     if first_is_false_target {
  64                         new_targets.all_targets_mut().swap(0, 1);
  65                     }
  66                 }
  67                 BinOp::Ne => {
  68                     // if the assignment was Ne we want the false case to be first
  69                     if !first_is_false_target {
  70                         new_targets.all_targets_mut().swap(0, 1);
  71                     }
  72                 }
  73                 _ => unreachable!(),
  74             }
  75
  76             // delete comparison statement if it the value being switched on was moved, which means it can not be user later on
  77             if opt.can_remove_bin_op_stmt {
  78                 bb.statements[opt.bin_op_stmt_idx].make_nop();
  79             } else {
  80                 // if the integer being compared to a const integral is being moved into the comparison,
  81                 // e.g `_2 = Eq(move _3, const 'x');`
  82                 // we want to avoid making a double move later on in the switchInt on _3.
  83                 // So to avoid `switchInt(move _3) -> ['x': bb2, otherwise: bb1];`,
  84                 // we convert the move in the comparison statement to a copy.
  85
  86                 // unwrap is safe as we know this statement is an assign
  87                 let (_, rhs) = bb.statements[opt.bin_op_stmt_idx].kind.as_assign_mut().unwrap();
  88
  89                 use Operand::*;
  90                 match rhs {
  91                     Rvalue::BinaryOp(_, box (ref mut left @ Move(_), Constant(_))) => {
  92                         *left = Copy(opt.to_switch_on);
  93                     }
  94                     Rvalue::BinaryOp(_, box (Constant(_), ref mut right @ Move(_))) => {
  95                         *right = Copy(opt.to_switch_on);
  96                     }
  97                     _ => (),
  98                 }
  99             }
 100
 101             let terminator = bb.terminator();
 102
 103             // remove StorageDead (if it exists) being used in the assign of the comparison
 104             for (stmt_idx, stmt) in bb.statements.iter().enumerate() {
 105                 if !matches!(stmt.kind, StatementKind::StorageDead(local) if local == opt.to_switch_on.local)
 106                 {
 107                     continue;
 108                 }
 109                 storage_deads_to_remove.push((stmt_idx, opt.bb_idx));
 110                 // if we have StorageDeads to remove then make sure to insert them at the top of each target
 111                 for bb_idx in new_targets.all_targets() {
 112                     storage_deads_to_insert.push((
 113                         *bb_idx,
 114                         Statement {
 115                             source_info: terminator.source_info,
 116                             kind: StatementKind::StorageDead(opt.to_switch_on.local),
 117                         },
 118                     ));
 119                 }
 120             }
 121
 122             let [bb_cond, bb_otherwise] = match new_targets.all_targets() {
 123                 [a, b] => [*a, *b],
 124                 e => bug!("expected 2 switch targets, got: {:?}", e),
 125             };
 126
 127             let targets = SwitchTargets::new(iter::once((new_value, bb_cond)), bb_otherwise);
 128
 129             let terminator = bb.terminator_mut();
 130             terminator.kind = TerminatorKind::SwitchInt {
 131                 discr: Operand::Move(opt.to_switch_on),
 132                 switch_ty: opt.branch_value_ty,
 133                 targets,
 134             };
 135         }
 136
 137         for (idx, bb_idx) in storage_deads_to_remove {
 138             body.basic_blocks_mut()[bb_idx].statements[idx].make_nop();
 139         }
 140
 141         for (idx, stmt) in storage_deads_to_insert {
 142             body.basic_blocks_mut()[idx].statements.insert(0, stmt);
 143         }
 144     }
 145 }
 146
 147 struct OptimizationFinder<'a, 'tcx> {
 148     body: &'a Body<'tcx>,
 149 }
 150
 151 impl<'tcx> OptimizationFinder<'_, 'tcx> {
 152     fn find_optimizations(&self) -> Vec<OptimizationInfo<'tcx>> {
 153         self.body
 154             .basic_blocks()
 155             .iter_enumerated()
 156             .filter_map(|(bb_idx, bb)| {
 157                 // find switch
 158                 let (place_switched_on, targets, place_switched_on_moved) =
 159                     match &bb.terminator().kind {
 160                         rustc_middle::mir::TerminatorKind::SwitchInt { discr, targets, .. } => {
 161                             Some((discr.place()?, targets, discr.is_move()))
 162                         }
 163                         _ => None,
 164                     }?;
 165
 166                 // find the statement that assigns the place being switched on
 167                 bb.statements.iter().enumerate().rev().find_map(|(stmt_idx, stmt)| {
 168                     match &stmt.kind {
 169                         rustc_middle::mir::StatementKind::Assign(box (lhs, rhs))
 170                             if *lhs == place_switched_on =>
 171                         {
 172                             match rhs {
 173                                 Rvalue::BinaryOp(
 174                                     op @ (BinOp::Eq | BinOp::Ne),
 175                                     box (left, right),
 176                                 ) => {
 177                                     let (branch_value_scalar, branch_value_ty, to_switch_on) =
 178                                         find_branch_value_info(left, right)?;
 179
 180                                     Some(OptimizationInfo {
 181                                         bin_op_stmt_idx: stmt_idx,
 182                                         bb_idx,
 183                                         can_remove_bin_op_stmt: place_switched_on_moved,
 184                                         to_switch_on,
 185                                         branch_value_scalar,
 186                                         branch_value_ty,
 187                                         op: *op,
 188                                         targets: targets.clone(),
 189                                     })
 190                                 }
 191                                 _ => None,
 192                             }
 193                         }
 194                         _ => None,
 195                     }
 196                 })
 197             })
 198             .collect()
 199     }
 200 }
 201
 202 fn find_branch_value_info<'tcx>(
 203     left: &Operand<'tcx>,
 204     right: &Operand<'tcx>,
 205 ) -> Option<(Scalar, Ty<'tcx>, Place<'tcx>)> {
 206     // check that either left or right is a constant.
 207     // if any are, we can use the other to switch on, and the constant as a value in a switch
 208     use Operand::*;
 209     match (left, right) {
 210         (Constant(branch_value), Copy(to_switch_on) | Move(to_switch_on))
 211         | (Copy(to_switch_on) | Move(to_switch_on), Constant(branch_value)) => {
 212             let branch_value_ty = branch_value.literal.ty();
 213             // we only want to apply this optimization if we are matching on integrals (and chars), as it is not possible to switch on floats
 214             if !branch_value_ty.is_integral() && !branch_value_ty.is_char() {
 215                 return None;
 216             };
 217             let branch_value_scalar = branch_value.literal.try_to_scalar()?;
 218             Some((branch_value_scalar, branch_value_ty, *to_switch_on))
 219         }
 220         _ => None,
 221     }
 222 }
 223
 224 #[derive(Debug)]
 225 struct OptimizationInfo<'tcx> {
 226     /// Basic block to apply the optimization
 227     bb_idx: BasicBlock,
 228     /// Statement index of Eq/Ne assignment that can be removed. None if the assignment can not be removed - i.e the statement is used later on
 229     bin_op_stmt_idx: usize,
 230     /// Can remove Eq/Ne assignment
 231     can_remove_bin_op_stmt: bool,
 232     /// Place that needs to be switched on. This place is of type integral
 233     to_switch_on: Place<'tcx>,
 234     /// Constant to use in switch target value
 235     branch_value_scalar: Scalar,
 236     /// Type of the constant value
 237     branch_value_ty: Ty<'tcx>,
 238     /// Either Eq or Ne
 239     op: BinOp,
 240     /// Current targets used in the switch
 241     targets: SwitchTargets,
 242 }