src/librustc_mir/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 // What's the intent of this pass?
   8 // If one block is found that switches between blocks which both go to the same place
   9 // AND both of these blocks set a similar const in their ->
  10 // condense into 1 block based on discriminant AND goto the destination afterwards
  11
  12 impl<'tcx> MirPass<'tcx> for MatchBranchSimplification {
  13     fn run_pass(&self, tcx: TyCtxt<'tcx>, src: MirSource<'tcx>, body: &mut Body<'tcx>) {
  14         let param_env = tcx.param_env(src.def_id());
  15         let bbs = body.basic_blocks_mut();
  16         'outer: for bb_idx in bbs.indices() {
  17             let (discr, val, switch_ty, first, second) = match bbs[bb_idx].terminator().kind {
  18                 TerminatorKind::SwitchInt {
  19                     discr: Operand::Copy(ref place) | Operand::Move(ref place),
  20                     switch_ty,
  21                     ref targets,
  22                     ref values,
  23                     ..
  24                 } if targets.len() == 2 && values.len() == 1 => {
  25                     (place, values[0], switch_ty, targets[0], targets[1])
  26                 }
  27                 // Only optimize switch int statements
  28                 _ => continue,
  29             };
  30
  31             // Check that destinations are identical, and if not, then don't optimize this block
  32             if &bbs[first].terminator().kind != &bbs[second].terminator().kind {
  33                 continue;
  34             }
  35
  36             // Check that blocks are assignments of consts to the same place or same statement,
  37             // and match up 1-1, if not don't optimize this block.
  38             let first_stmts = &bbs[first].statements;
  39             let scnd_stmts = &bbs[second].statements;
  40             if first_stmts.len() != scnd_stmts.len() {
  41                 continue;
  42             }
  43             for (f, s) in first_stmts.iter().zip(scnd_stmts.iter()) {
  44                 match (&f.kind, &s.kind) {
  45                     // If two statements are exactly the same just ignore them.
  46                     (f_s, s_s) if f_s == s_s => (),
  47
  48                     (
  49                         StatementKind::Assign(box (lhs_f, Rvalue::Use(Operand::Constant(f_c)))),
  50                         StatementKind::Assign(box (lhs_s, Rvalue::Use(Operand::Constant(s_c)))),
  51                     ) if lhs_f == lhs_s && f_c.literal.ty.is_bool() && s_c.literal.ty.is_bool() => {
  52                         let f_c = f_c.literal.try_eval_bool(tcx, param_env).unwrap();
  53                         let s_c = s_c.literal.try_eval_bool(tcx, param_env).unwrap();
  54                         if f_c != s_c {
  55                             // have to check this here because f_c & s_c might have
  56                             // different spans.
  57                             continue;
  58                         }
  59                         continue 'outer;
  60                     }
  61                     // If there are not exclusively assignments, then ignore this
  62                     _ => continue 'outer,
  63                 }
  64             }
  65             // Take owenership of items now that we know we can optimize.
  66             let discr = discr.clone();
  67             let (from, first) = bbs.pick2_mut(bb_idx, first);
  68
  69             let new_stmts = first.statements.iter().cloned().map(|mut s| {
  70                 if let StatementKind::Assign(box (_, ref mut rhs)) = s.kind {
  71                     if let Rvalue::Use(Operand::Constant(c)) = rhs {
  72                         let size = tcx.layout_of(param_env.and(switch_ty)).unwrap().size;
  73                         let const_cmp = Operand::const_from_scalar(
  74                             tcx,
  75                             switch_ty,
  76                             crate::interpret::Scalar::from_uint(val, size),
  77                             rustc_span::DUMMY_SP,
  78                         );
  79                         if let Some(c) = c.literal.try_eval_bool(tcx, param_env) {
  80                             let op = if c { BinOp::Eq } else { BinOp::Ne };
  81                             *rhs = Rvalue::BinaryOp(op, Operand::Copy(discr), const_cmp);
  82                         }
  83                     }
  84                 }
  85                 s
  86             });
  87             from.statements.extend(new_stmts);
  88             from.terminator_mut().kind = first.terminator().kind.clone();
  89         }
  90     }
  91 }