compiler/rustc_mir_transform/src/uninhabited_enum_branching.rs

   1 //! A pass that eliminates branches on uninhabited enum variants.
   2
   3 use crate::MirPass;
   4 use rustc_data_structures::fx::FxHashSet;
   5 use rustc_middle::mir::{
   6     BasicBlockData, Body, Local, Operand, Rvalue, StatementKind, SwitchTargets, Terminator,
   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     tcx: TyCtxt<'tcx>,
  28     body: &Body<'tcx>,
  29 ) -> Option<Ty<'tcx>> {
  30     let terminator = block_data.terminator();
  31
  32     // Only bother checking blocks which terminate by switching on a local.
  33     if let Some(local) = get_discriminant_local(&terminator.kind) {
  34         let stmt_before_term = (!block_data.statements.is_empty())
  35             .then(|| &block_data.statements[block_data.statements.len() - 1].kind);
  36
  37         if let Some(StatementKind::Assign(box (l, Rvalue::Discriminant(place)))) = stmt_before_term
  38         {
  39             if l.as_local() == Some(local) {
  40                 let ty = place.ty(body, tcx).ty;
  41                 if ty.is_enum() {
  42                     return Some(ty);
  43                 }
  44             }
  45         }
  46     }
  47
  48     None
  49 }
  50
  51 fn variant_discriminants<'tcx>(
  52     layout: &TyAndLayout<'tcx>,
  53     ty: Ty<'tcx>,
  54     tcx: TyCtxt<'tcx>,
  55 ) -> FxHashSet<u128> {
  56     match &layout.variants {
  57         Variants::Single { index } => {
  58             let mut res = FxHashSet::default();
  59             res.insert(
  60                 ty.discriminant_for_variant(tcx, *index)
  61                     .map_or(index.as_u32() as u128, |discr| discr.val),
  62             );
  63             res
  64         }
  65         Variants::Multiple { variants, .. } => variants
  66             .iter_enumerated()
  67             .filter_map(|(idx, layout)| {
  68                 (layout.abi != Abi::Uninhabited)
  69                     .then(|| ty.discriminant_for_variant(tcx, idx).unwrap().val)
  70             })
  71             .collect(),
  72     }
  73 }
  74
  75 /// Ensures that the `otherwise` branch leads to an unreachable bb, returning `None` if so and a new
  76 /// bb to use as the new target if not.
  77 fn ensure_otherwise_unreachable<'tcx>(
  78     body: &Body<'tcx>,
  79     targets: &SwitchTargets,
  80 ) -> Option<BasicBlockData<'tcx>> {
  81     let otherwise = targets.otherwise();
  82     let bb = &body.basic_blocks[otherwise];
  83     if bb.terminator().kind == TerminatorKind::Unreachable
  84         && bb.statements.iter().all(|s| matches!(&s.kind, StatementKind::StorageDead(_)))
  85     {
  86         return None;
  87     }
  88
  89     let mut new_block = BasicBlockData::new(Some(Terminator {
  90         source_info: bb.terminator().source_info,
  91         kind: TerminatorKind::Unreachable,
  92     }));
  93     new_block.is_cleanup = bb.is_cleanup;
  94     Some(new_block)
  95 }
  96
  97 impl<'tcx> MirPass<'tcx> for UninhabitedEnumBranching {
  98     fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
  99         sess.mir_opt_level() > 0
 100     }
 101
 102     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
 103         trace!("UninhabitedEnumBranching starting for {:?}", body.source);
 104
 105         for bb in body.basic_blocks.indices() {
 106             trace!("processing block {:?}", bb);
 107
 108             let Some(discriminant_ty) = get_switched_on_type(&body.basic_blocks[bb], tcx, body) else {
 109                 continue;
 110             };
 111
 112             let layout = tcx.layout_of(tcx.param_env(body.source.def_id()).and(discriminant_ty));
 113
 114             let allowed_variants = if let Ok(layout) = layout {
 115                 variant_discriminants(&layout, discriminant_ty, tcx)
 116             } else {
 117                 continue;
 118             };
 119
 120             trace!("allowed_variants = {:?}", allowed_variants);
 121
 122             if let TerminatorKind::SwitchInt { targets, .. } =
 123                 &mut body.basic_blocks_mut()[bb].terminator_mut().kind
 124             {
 125                 let mut new_targets = SwitchTargets::new(
 126                     targets.iter().filter(|(val, _)| allowed_variants.contains(val)),
 127                     targets.otherwise(),
 128                 );
 129
 130                 if new_targets.iter().count() == allowed_variants.len() {
 131                     if let Some(updated) = ensure_otherwise_unreachable(body, &new_targets) {
 132                         let new_otherwise = body.basic_blocks_mut().push(updated);
 133                         *new_targets.all_targets_mut().last_mut().unwrap() = new_otherwise;
 134                     }
 135                 }
 136
 137                 if let TerminatorKind::SwitchInt { targets, .. } =
 138                     &mut body.basic_blocks_mut()[bb].terminator_mut().kind
 139                 {
 140                     *targets = new_targets;
 141                 } else {
 142                     unreachable!()
 143                 }
 144             } else {
 145                 unreachable!()
 146             }
 147         }
 148     }
 149 }