1 //! A pass that eliminates branches on uninhabited enum variants.
4 use rustc_data_structures::fx::FxHashSet;
5 use rustc_middle::mir::{
6 BasicBlockData, Body, Local, Operand, Rvalue, StatementKind, SwitchTargets, Terminator,
9 use rustc_middle::ty::layout::TyAndLayout;
10 use rustc_middle::ty::{Ty, TyCtxt};
11 use rustc_target::abi::{Abi, Variants};
13 pub struct UninhabitedEnumBranching;
15 fn get_discriminant_local(terminator: &TerminatorKind<'_>) -> Option<Local> {
16 if let TerminatorKind::SwitchInt { discr: Operand::Move(p), .. } = terminator {
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>,
29 ) -> Option<Ty<'tcx>> {
30 let terminator = block_data.terminator();
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);
37 if let Some(StatementKind::Assign(box (l, Rvalue::Discriminant(place)))) = stmt_before_term
39 if l.as_local() == Some(local) {
40 let ty = place.ty(body, tcx).ty;
51 fn variant_discriminants<'tcx>(
52 layout: &TyAndLayout<'tcx>,
55 ) -> FxHashSet<u128> {
56 match &layout.variants {
57 Variants::Single { index } => {
58 let mut res = FxHashSet::default();
60 ty.discriminant_for_variant(tcx, *index)
61 .map_or(index.as_u32() as u128, |discr| discr.val),
65 Variants::Multiple { variants, .. } => variants
67 .filter_map(|(idx, layout)| {
68 (layout.abi() != Abi::Uninhabited)
69 .then(|| ty.discriminant_for_variant(tcx, idx).unwrap().val)
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>(
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(_)))
89 let mut new_block = BasicBlockData::new(Some(Terminator {
90 source_info: bb.terminator().source_info,
91 kind: TerminatorKind::Unreachable,
93 new_block.is_cleanup = bb.is_cleanup;
97 impl<'tcx> MirPass<'tcx> for UninhabitedEnumBranching {
98 fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
99 sess.mir_opt_level() > 0
102 fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
103 trace!("UninhabitedEnumBranching starting for {:?}", body.source);
105 for bb in body.basic_blocks.indices() {
106 trace!("processing block {:?}", bb);
108 let Some(discriminant_ty) = get_switched_on_type(&body.basic_blocks[bb], tcx, body) else {
112 let layout = tcx.layout_of(tcx.param_env(body.source.def_id()).and(discriminant_ty));
114 let allowed_variants = if let Ok(layout) = layout {
115 variant_discriminants(&layout, discriminant_ty, tcx)
120 trace!("allowed_variants = {:?}", allowed_variants);
122 if let TerminatorKind::SwitchInt { targets, .. } =
123 &mut body.basic_blocks_mut()[bb].terminator_mut().kind
125 let mut new_targets = SwitchTargets::new(
126 targets.iter().filter(|(val, _)| allowed_variants.contains(val)),
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;
137 if let TerminatorKind::SwitchInt { targets, .. } =
138 &mut body.basic_blocks_mut()[bb].terminator_mut().kind
140 *targets = new_targets;