compiler/rustc_mir/src/transform/const_goto.rs

   1 //! This pass optimizes the following sequence
   2 //! ```rust,ignore (example)
   3 //! bb2: {
   4 //!     _2 = const true;
   5 //!     goto -> bb3;
   6 //! }
   7 //!
   8 //! bb3: {
   9 //!     switchInt(_2) -> [false: bb4, otherwise: bb5];
  10 //! }
  11 //! ```
  12 //! into
  13 //! ```rust,ignore (example)
  14 //! bb2: {
  15 //!     _2 = const true;
  16 //!     goto -> bb5;
  17 //! }
  18 //! ```
  19
  20 use crate::transform::MirPass;
  21 use rustc_middle::mir::*;
  22 use rustc_middle::ty::TyCtxt;
  23 use rustc_middle::{mir::visit::Visitor, ty::ParamEnv};
  24
  25 use super::simplify::{simplify_cfg, simplify_locals};
  26
  27 pub struct ConstGoto;
  28
  29 impl<'tcx> MirPass<'tcx> for ConstGoto {
  30     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
  31         if tcx.sess.mir_opt_level() < 4 {
  32             return;
  33         }
  34         trace!("Running ConstGoto on {:?}", body.source);
  35         let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id());
  36         let mut opt_finder =
  37             ConstGotoOptimizationFinder { tcx, body, optimizations: vec![], param_env };
  38         opt_finder.visit_body(body);
  39         let should_simplify = !opt_finder.optimizations.is_empty();
  40         for opt in opt_finder.optimizations {
  41             let terminator = body.basic_blocks_mut()[opt.bb_with_goto].terminator_mut();
  42             let new_goto = TerminatorKind::Goto { target: opt.target_to_use_in_goto };
  43             debug!("SUCCESS: replacing `{:?}` with `{:?}`", terminator.kind, new_goto);
  44             terminator.kind = new_goto;
  45         }
  46
  47         // if we applied optimizations, we potentially have some cfg to cleanup to
  48         // make it easier for further passes
  49         if should_simplify {
  50             simplify_cfg(body);
  51             simplify_locals(body, tcx);
  52         }
  53     }
  54 }
  55
  56 impl<'a, 'tcx> Visitor<'tcx> for ConstGotoOptimizationFinder<'a, 'tcx> {
  57     fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) {
  58         let _: Option<_> = try {
  59             let target = terminator.kind.as_goto()?;
  60             // We only apply this optimization if the last statement is a const assignment
  61             let last_statement = self.body.basic_blocks()[location.block].statements.last()?;
  62
  63             if let (place, Rvalue::Use(Operand::Constant(_const))) =
  64                 last_statement.kind.as_assign()?
  65             {
  66                 // We found a constant being assigned to `place`.
  67                 // Now check that the target of this Goto switches on this place.
  68                 let target_bb = &self.body.basic_blocks()[target];
  69
  70                 // FIXME(simonvandel): We are conservative here when we don't allow
  71                 // any statements in the target basic block.
  72                 // This could probably be relaxed to allow `StorageDead`s which could be
  73                 // copied to the predecessor of this block.
  74                 if !target_bb.statements.is_empty() {
  75                     None?
  76                 }
  77
  78                 let target_bb_terminator = target_bb.terminator();
  79                 let (discr, switch_ty, targets) = target_bb_terminator.kind.as_switch()?;
  80                 if discr.place() == Some(*place) {
  81                     // We now know that the Switch matches on the const place, and it is statementless
  82                     // Now find which value in the Switch matches the const value.
  83                     let const_value =
  84                         _const.literal.try_eval_bits(self.tcx, self.param_env, switch_ty)?;
  85                     let found_value_idx_option = targets
  86                         .iter()
  87                         .enumerate()
  88                         .find(|(_, (value, _))| const_value == *value)
  89                         .map(|(idx, _)| idx);
  90
  91                     let target_to_use_in_goto =
  92                         if let Some(found_value_idx) = found_value_idx_option {
  93                             targets.iter().nth(found_value_idx).unwrap().1
  94                         } else {
  95                             // If we did not find the const value in values, it must be the otherwise case
  96                             targets.otherwise()
  97                         };
  98
  99                     self.optimizations.push(OptimizationToApply {
 100                         bb_with_goto: location.block,
 101                         target_to_use_in_goto,
 102                     });
 103                 }
 104             }
 105             Some(())
 106         };
 107
 108         self.super_terminator(terminator, location);
 109     }
 110 }
 111
 112 struct OptimizationToApply {
 113     bb_with_goto: BasicBlock,
 114     target_to_use_in_goto: BasicBlock,
 115 }
 116
 117 pub struct ConstGotoOptimizationFinder<'a, 'tcx> {
 118     tcx: TyCtxt<'tcx>,
 119     body: &'a Body<'tcx>,
 120     param_env: ParamEnv<'tcx>,
 121     optimizations: Vec<OptimizationToApply>,
 122 }