| //! A pass that eliminates branches on uninhabited enum variants. |
| |
| use crate::MirPass; |
| use rustc_data_structures::fx::FxHashSet; |
| use rustc_middle::mir::{ |
| BasicBlockData, Body, Local, Operand, Rvalue, StatementKind, SwitchTargets, Terminator, |
| TerminatorKind, |
| }; |
| use rustc_middle::ty::layout::TyAndLayout; |
| use rustc_middle::ty::{Ty, TyCtxt}; |
| use rustc_target::abi::{Abi, Variants}; |
| |
| pub struct UninhabitedEnumBranching; |
| |
| fn get_discriminant_local(terminator: &TerminatorKind<'_>) -> Option<Local> { |
| if let TerminatorKind::SwitchInt { discr: Operand::Move(p), .. } = terminator { |
| p.as_local() |
| } else { |
| None |
| } |
| } |
| |
| /// If the basic block terminates by switching on a discriminant, this returns the `Ty` the |
| /// discriminant is read from. Otherwise, returns None. |
| fn get_switched_on_type<'tcx>( |
| block_data: &BasicBlockData<'tcx>, |
| tcx: TyCtxt<'tcx>, |
| body: &Body<'tcx>, |
| ) -> Option<Ty<'tcx>> { |
| let terminator = block_data.terminator(); |
| |
| // Only bother checking blocks which terminate by switching on a local. |
| if let Some(local) = get_discriminant_local(&terminator.kind) { |
| let stmt_before_term = (!block_data.statements.is_empty()) |
| .then(|| &block_data.statements[block_data.statements.len() - 1].kind); |
| |
| if let Some(StatementKind::Assign(box (l, Rvalue::Discriminant(place)))) = stmt_before_term |
| { |
| if l.as_local() == Some(local) { |
| let ty = place.ty(body, tcx).ty; |
| if ty.is_enum() { |
| return Some(ty); |
| } |
| } |
| } |
| } |
| |
| None |
| } |
| |
| fn variant_discriminants<'tcx>( |
| layout: &TyAndLayout<'tcx>, |
| ty: Ty<'tcx>, |
| tcx: TyCtxt<'tcx>, |
| ) -> FxHashSet<u128> { |
| match &layout.variants { |
| Variants::Single { index } => { |
| let mut res = FxHashSet::default(); |
| res.insert( |
| ty.discriminant_for_variant(tcx, *index) |
| .map_or(index.as_u32() as u128, |discr| discr.val), |
| ); |
| res |
| } |
| Variants::Multiple { variants, .. } => variants |
| .iter_enumerated() |
| .filter_map(|(idx, layout)| { |
| (layout.abi != Abi::Uninhabited) |
| .then(|| ty.discriminant_for_variant(tcx, idx).unwrap().val) |
| }) |
| .collect(), |
| } |
| } |
| |
| /// Ensures that the `otherwise` branch leads to an unreachable bb, returning `None` if so and a new |
| /// bb to use as the new target if not. |
| fn ensure_otherwise_unreachable<'tcx>( |
| body: &Body<'tcx>, |
| targets: &SwitchTargets, |
| ) -> Option<BasicBlockData<'tcx>> { |
| let otherwise = targets.otherwise(); |
| let bb = &body.basic_blocks[otherwise]; |
| if bb.terminator().kind == TerminatorKind::Unreachable |
| && bb.statements.iter().all(|s| matches!(&s.kind, StatementKind::StorageDead(_))) |
| { |
| return None; |
| } |
| |
| let mut new_block = BasicBlockData::new(Some(Terminator { |
| source_info: bb.terminator().source_info, |
| kind: TerminatorKind::Unreachable, |
| })); |
| new_block.is_cleanup = bb.is_cleanup; |
| Some(new_block) |
| } |
| |
| impl<'tcx> MirPass<'tcx> for UninhabitedEnumBranching { |
| fn is_enabled(&self, sess: &rustc_session::Session) -> bool { |
| sess.mir_opt_level() > 0 |
| } |
| |
| fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
| trace!("UninhabitedEnumBranching starting for {:?}", body.source); |
| |
| for bb in body.basic_blocks.indices() { |
| trace!("processing block {:?}", bb); |
| |
| let Some(discriminant_ty) = get_switched_on_type(&body.basic_blocks[bb], tcx, body) |
| else { |
| continue; |
| }; |
| |
| let layout = tcx.layout_of( |
| tcx.param_env_reveal_all_normalized(body.source.def_id()).and(discriminant_ty), |
| ); |
| |
| let allowed_variants = if let Ok(layout) = layout { |
| variant_discriminants(&layout, discriminant_ty, tcx) |
| } else { |
| continue; |
| }; |
| |
| trace!("allowed_variants = {:?}", allowed_variants); |
| |
| if let TerminatorKind::SwitchInt { targets, .. } = |
| &mut body.basic_blocks_mut()[bb].terminator_mut().kind |
| { |
| let mut new_targets = SwitchTargets::new( |
| targets.iter().filter(|(val, _)| allowed_variants.contains(val)), |
| targets.otherwise(), |
| ); |
| |
| if new_targets.iter().count() == allowed_variants.len() { |
| if let Some(updated) = ensure_otherwise_unreachable(body, &new_targets) { |
| let new_otherwise = body.basic_blocks_mut().push(updated); |
| *new_targets.all_targets_mut().last_mut().unwrap() = new_otherwise; |
| } |
| } |
| |
| if let TerminatorKind::SwitchInt { targets, .. } = |
| &mut body.basic_blocks_mut()[bb].terminator_mut().kind |
| { |
| *targets = new_targets; |
| } else { |
| unreachable!() |
| } |
| } else { |
| unreachable!() |
| } |
| } |
| } |
| } |