compiler/rustc_mir/src/transform/match_branches.rs

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