compiler/rustc_mir/src/transform/simplify_comparison_integral.rs

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