compiler/rustc_mir_transform/src/match_branches.rs

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