compiler/rustc_mir/src/transform/match_branches.rs

   1 use crate::transform::MirPass;
   2 use rustc_middle::mir::*;
   3 use rustc_middle::ty::TyCtxt;
   4
   5 use super::simplify::simplify_cfg;
   6
   7 pub struct MatchBranchSimplification;
   8
   9 /// If a source block is found that switches between two blocks that are exactly
  10 /// the same modulo const bool assignments (e.g., one assigns true another false
  11 /// to the same place), merge a target block statements into the source block,
  12 /// using Eq / Ne comparison with switch value where const bools value differ.
  13 ///
  14 /// For example:
  15 ///
  16 /// ```rust
  17 /// bb0: {
  18 ///     switchInt(move _3) -> [42_isize: bb1, otherwise: bb2];
  19 /// }
  20 ///
  21 /// bb1: {
  22 ///     _2 = const true;
  23 ///     goto -> bb3;
  24 /// }
  25 ///
  26 /// bb2: {
  27 ///     _2 = const false;
  28 ///     goto -> bb3;
  29 /// }
  30 /// ```
  31 ///
  32 /// into:
  33 ///
  34 /// ```rust
  35 /// bb0: {
  36 ///    _2 = Eq(move _3, const 42_isize);
  37 ///    goto -> bb3;
  38 /// }
  39 /// ```
  40
  41 impl<'tcx> MirPass<'tcx> for MatchBranchSimplification {
  42     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
  43         if tcx.sess.mir_opt_level() < 3 {
  44             return;
  45         }
  46
  47         let def_id = body.source.def_id();
  48         let param_env = tcx.param_env(def_id);
  49
  50         let (bbs, local_decls) = body.basic_blocks_and_local_decls_mut();
  51         let mut should_cleanup = false;
  52         'outer: for bb_idx in bbs.indices() {
  53             if !tcx.consider_optimizing(|| format!("MatchBranchSimplification {:?} ", def_id)) {
  54                 continue;
  55             }
  56
  57             let (discr, val, switch_ty, first, second) = match bbs[bb_idx].terminator().kind {
  58                 TerminatorKind::SwitchInt {
  59                     discr: ref discr @ (Operand::Copy(_) | Operand::Move(_)),
  60                     switch_ty,
  61                     ref targets,
  62                     ..
  63                 } if targets.iter().len() == 1 => {
  64                     let (value, target) = targets.iter().next().unwrap();
  65                     if target == targets.otherwise() {
  66                         continue;
  67                     }
  68                     (discr, value, switch_ty, target, targets.otherwise())
  69                 }
  70                 // Only optimize switch int statements
  71                 _ => continue,
  72             };
  73
  74             // Check that destinations are identical, and if not, then don't optimize this block
  75             if bbs[first].terminator().kind != bbs[second].terminator().kind {
  76                 continue;
  77             }
  78
  79             // Check that blocks are assignments of consts to the same place or same statement,
  80             // and match up 1-1, if not don't optimize this block.
  81             let first_stmts = &bbs[first].statements;
  82             let scnd_stmts = &bbs[second].statements;
  83             if first_stmts.len() != scnd_stmts.len() {
  84                 continue;
  85             }
  86             for (f, s) in first_stmts.iter().zip(scnd_stmts.iter()) {
  87                 match (&f.kind, &s.kind) {
  88                     // If two statements are exactly the same, we can optimize.
  89                     (f_s, s_s) if f_s == s_s => {}
  90
  91                     // If two statements are const bool assignments to the same place, we can optimize.
  92                     (
  93                         StatementKind::Assign(box (lhs_f, Rvalue::Use(Operand::Constant(f_c)))),
  94                         StatementKind::Assign(box (lhs_s, Rvalue::Use(Operand::Constant(s_c)))),
  95                     ) if lhs_f == lhs_s
  96                         && f_c.literal.ty.is_bool()
  97                         && s_c.literal.ty.is_bool()
  98                         && f_c.literal.try_eval_bool(tcx, param_env).is_some()
  99                         && s_c.literal.try_eval_bool(tcx, param_env).is_some() => {}
 100
 101                     // Otherwise we cannot optimize. Try another block.
 102                     _ => continue 'outer,
 103                 }
 104             }
 105             // Take ownership of items now that we know we can optimize.
 106             let discr = discr.clone();
 107
 108             // Introduce a temporary for the discriminant value.
 109             let source_info = bbs[bb_idx].terminator().source_info;
 110             let discr_local = local_decls.push(LocalDecl::new(switch_ty, source_info.span));
 111
 112             // We already checked that first and second are different blocks,
 113             // and bb_idx has a different terminator from both of them.
 114             let (from, first, second) = bbs.pick3_mut(bb_idx, first, second);
 115
 116             let new_stmts = first.statements.iter().zip(second.statements.iter()).map(|(f, s)| {
 117                 match (&f.kind, &s.kind) {
 118                     (f_s, s_s) if f_s == s_s => (*f).clone(),
 119
 120                     (
 121                         StatementKind::Assign(box (lhs, Rvalue::Use(Operand::Constant(f_c)))),
 122                         StatementKind::Assign(box (_, Rvalue::Use(Operand::Constant(s_c)))),
 123                     ) => {
 124                         // From earlier loop we know that we are dealing with bool constants only:
 125                         let f_b = f_c.literal.try_eval_bool(tcx, param_env).unwrap();
 126                         let s_b = s_c.literal.try_eval_bool(tcx, param_env).unwrap();
 127                         if f_b == s_b {
 128                             // Same value in both blocks. Use statement as is.
 129                             (*f).clone()
 130                         } else {
 131                             // Different value between blocks. Make value conditional on switch condition.
 132                             let size = tcx.layout_of(param_env.and(switch_ty)).unwrap().size;
 133                             let const_cmp = Operand::const_from_scalar(
 134                                 tcx,
 135                                 switch_ty,
 136                                 crate::interpret::Scalar::from_uint(val, size),
 137                                 rustc_span::DUMMY_SP,
 138                             );
 139                             let op = if f_b { BinOp::Eq } else { BinOp::Ne };
 140                             let rhs = Rvalue::BinaryOp(
 141                                 op,
 142                                 Operand::Copy(Place::from(discr_local)),
 143                                 const_cmp,
 144                             );
 145                             Statement {
 146                                 source_info: f.source_info,
 147                                 kind: StatementKind::Assign(box (*lhs, rhs)),
 148                             }
 149                         }
 150                     }
 151
 152                     _ => unreachable!(),
 153                 }
 154             });
 155
 156             from.statements
 157                 .push(Statement { source_info, kind: StatementKind::StorageLive(discr_local) });
 158             from.statements.push(Statement {
 159                 source_info,
 160                 kind: StatementKind::Assign(box (Place::from(discr_local), Rvalue::Use(discr))),
 161             });
 162             from.statements.extend(new_stmts);
 163             from.statements
 164                 .push(Statement { source_info, kind: StatementKind::StorageDead(discr_local) });
 165             from.terminator_mut().kind = first.terminator().kind.clone();
 166             should_cleanup = true;
 167         }
 168
 169         if should_cleanup {
 170             simplify_cfg(body);
 171         }
 172     }
 173 }