compiler/rustc_mir/src/transform/uninhabited_enum_branching.rs

   1 //! A pass that eliminates branches on uninhabited enum variants.
   2
   3 use crate::transform::MirPass;
   4 use rustc_data_structures::stable_set::FxHashSet;
   5 use rustc_middle::mir::{
   6     BasicBlock, BasicBlockData, Body, Local, Operand, Rvalue, StatementKind, SwitchTargets,
   7     TerminatorKind,
   8 };
   9 use rustc_middle::ty::layout::TyAndLayout;
  10 use rustc_middle::ty::{Ty, TyCtxt};
  11 use rustc_target::abi::{Abi, Variants};
  12
  13 pub struct UninhabitedEnumBranching;
  14
  15 fn get_discriminant_local(terminator: &TerminatorKind<'_>) -> Option<Local> {
  16     if let TerminatorKind::SwitchInt { discr: Operand::Move(p), .. } = terminator {
  17         p.as_local()
  18     } else {
  19         None
  20     }
  21 }
  22
  23 /// If the basic block terminates by switching on a discriminant, this returns the `Ty` the
  24 /// discriminant is read from. Otherwise, returns None.
  25 fn get_switched_on_type<'tcx>(
  26     block_data: &BasicBlockData<'tcx>,
  27     body: &Body<'tcx>,
  28 ) -> Option<Ty<'tcx>> {
  29     let terminator = block_data.terminator();
  30
  31     // Only bother checking blocks which terminate by switching on a local.
  32     if let Some(local) = get_discriminant_local(&terminator.kind) {
  33         let stmt_before_term = (!block_data.statements.is_empty())
  34             .then(|| &block_data.statements[block_data.statements.len() - 1].kind);
  35
  36         if let Some(StatementKind::Assign(box (l, Rvalue::Discriminant(place)))) = stmt_before_term
  37         {
  38             if l.as_local() == Some(local) {
  39                 if let Some(r_local) = place.as_local() {
  40                     let ty = body.local_decls[r_local].ty;
  41
  42                     if ty.is_enum() {
  43                         return Some(ty);
  44                     }
  45                 }
  46             }
  47         }
  48     }
  49
  50     None
  51 }
  52
  53 fn variant_discriminants<'tcx>(
  54     layout: &TyAndLayout<'tcx>,
  55     ty: Ty<'tcx>,
  56     tcx: TyCtxt<'tcx>,
  57 ) -> FxHashSet<u128> {
  58     match &layout.variants {
  59         Variants::Single { index } => {
  60             let mut res = FxHashSet::default();
  61             res.insert(index.as_u32() as u128);
  62             res
  63         }
  64         Variants::Multiple { variants, .. } => variants
  65             .iter_enumerated()
  66             .filter_map(|(idx, layout)| {
  67                 (layout.abi != Abi::Uninhabited)
  68                     .then(|| ty.discriminant_for_variant(tcx, idx).unwrap().val)
  69             })
  70             .collect(),
  71     }
  72 }
  73
  74 impl<'tcx> MirPass<'tcx> for UninhabitedEnumBranching {
  75     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
  76         if body.source.promoted.is_some() {
  77             return;
  78         }
  79
  80         trace!("UninhabitedEnumBranching starting for {:?}", body.source);
  81
  82         let basic_block_count = body.basic_blocks().len();
  83
  84         for bb in 0..basic_block_count {
  85             let bb = BasicBlock::from_usize(bb);
  86             trace!("processing block {:?}", bb);
  87
  88             let discriminant_ty =
  89                 if let Some(ty) = get_switched_on_type(&body.basic_blocks()[bb], body) {
  90                     ty
  91                 } else {
  92                     continue;
  93                 };
  94
  95             let layout = tcx.layout_of(tcx.param_env(body.source.def_id()).and(discriminant_ty));
  96
  97             let allowed_variants = if let Ok(layout) = layout {
  98                 variant_discriminants(&layout, discriminant_ty, tcx)
  99             } else {
 100                 continue;
 101             };
 102
 103             trace!("allowed_variants = {:?}", allowed_variants);
 104
 105             if let TerminatorKind::SwitchInt { targets, .. } =
 106                 &mut body.basic_blocks_mut()[bb].terminator_mut().kind
 107             {
 108                 let new_targets = SwitchTargets::new(
 109                     targets.iter().filter(|(val, _)| allowed_variants.contains(val)),
 110                     targets.otherwise(),
 111                 );
 112
 113                 *targets = new_targets;
 114             } else {
 115                 unreachable!()
 116             }
 117         }
 118     }
 119 }