compiler/rustc_mir_transform/src/simplify_branches.rs

   1 use crate::MirPass;
   2 use rustc_middle::mir::*;
   3 use rustc_middle::ty::TyCtxt;
   4
   5 use std::borrow::Cow;
   6
   7 /// A pass that replaces a branch with a goto when its condition is known.
   8 pub struct SimplifyConstCondition {
   9     label: String,
  10 }
  11
  12 impl SimplifyConstCondition {
  13     pub fn new(label: &str) -> Self {
  14         SimplifyConstCondition { label: format!("SimplifyConstCondition-{}", label) }
  15     }
  16 }
  17
  18 impl<'tcx> MirPass<'tcx> for SimplifyConstCondition {
  19     fn name(&self) -> Cow<'_, str> {
  20         Cow::Borrowed(&self.label)
  21     }
  22
  23     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
  24         let param_env = tcx.param_env(body.source.def_id());
  25         for block in body.basic_blocks_mut() {
  26             let terminator = block.terminator_mut();
  27             terminator.kind = match terminator.kind {
  28                 TerminatorKind::SwitchInt {
  29                     discr: Operand::Constant(ref c),
  30                     switch_ty,
  31                     ref targets,
  32                     ..
  33                 } => {
  34                     let constant = c.literal.try_eval_bits(tcx, param_env, switch_ty);
  35                     if let Some(constant) = constant {
  36                         let target = targets.target_for_value(constant);
  37                         TerminatorKind::Goto { target }
  38                     } else {
  39                         continue;
  40                     }
  41                 }
  42                 TerminatorKind::Assert {
  43                     target, cond: Operand::Constant(ref c), expected, ..
  44                 } => match c.literal.try_eval_bool(tcx, param_env) {
  45                     Some(v) if v == expected => TerminatorKind::Goto { target },
  46                     _ => continue,
  47                 },
  48                 _ => continue,
  49             };
  50         }
  51     }
  52 }