1 //! A pass that eliminates branches on uninhabited enum variants.
2
3 use crate::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
get_discriminant_local(terminator: &TerminatorKind<'_>) -> Option<Local>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.
get_switched_on_type<'tcx>( block_data: &BasicBlockData<'tcx>, tcx: TyCtxt<'tcx>, body: &Body<'tcx>, ) -> Option<Ty<'tcx>>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
variant_discriminants<'tcx>( layout: &TyAndLayout<'tcx>, ty: Ty<'tcx>, tcx: TyCtxt<'tcx>, ) -> FxHashSet<u128>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(index.as_u32() as u128);
60 res
61 }
62 Variants::Multiple { variants, .. } => variants
63 .iter_enumerated()
64 .filter_map(|(idx, layout)| {
65 (layout.abi != Abi::Uninhabited)
66 .then(|| ty.discriminant_for_variant(tcx, idx).unwrap().val)
67 })
68 .collect(),
69 }
70 }
71
72 impl<'tcx> MirPass<'tcx> for UninhabitedEnumBranching {
run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>)73 fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
74 if body.source.promoted.is_some() {
75 return;
76 }
77
78 trace!("UninhabitedEnumBranching starting for {:?}", body.source);
79
80 let basic_block_count = body.basic_blocks().len();
81
82 for bb in 0..basic_block_count {
83 let bb = BasicBlock::from_usize(bb);
84 trace!("processing block {:?}", bb);
85
86 let Some(discriminant_ty) = get_switched_on_type(&body.basic_blocks()[bb], tcx, body) else {
87 continue;
88 };
89
90 let layout = tcx.layout_of(tcx.param_env(body.source.def_id()).and(discriminant_ty));
91
92 let allowed_variants = if let Ok(layout) = layout {
93 variant_discriminants(&layout, discriminant_ty, tcx)
94 } else {
95 continue;
96 };
97
98 trace!("allowed_variants = {:?}", allowed_variants);
99
100 if let TerminatorKind::SwitchInt { targets, .. } =
101 &mut body.basic_blocks_mut()[bb].terminator_mut().kind
102 {
103 let new_targets = SwitchTargets::new(
104 targets.iter().filter(|(val, _)| allowed_variants.contains(val)),
105 targets.otherwise(),
106 );
107
108 *targets = new_targets;
109 } else {
110 unreachable!()
111 }
112 }
113 }
114 }
115