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