]>
Commit | Line | Data |
---|---|---|
60c5eb7d XL |
1 | //! A pass that eliminates branches on uninhabited enum variants. |
2 | ||
c295e0f8 | 3 | use crate::MirPass; |
29967ef6 | 4 | use rustc_data_structures::stable_set::FxHashSet; |
ba9703b0 | 5 | use rustc_middle::mir::{ |
5e7ed085 FG |
6 | BasicBlockData, Body, Local, Operand, Rvalue, StatementKind, SwitchTargets, Terminator, |
7 | TerminatorKind, | |
60c5eb7d | 8 | }; |
ba9703b0 XL |
9 | use rustc_middle::ty::layout::TyAndLayout; |
10 | use rustc_middle::ty::{Ty, TyCtxt}; | |
11 | use rustc_target::abi::{Abi, Variants}; | |
60c5eb7d XL |
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>, | |
136023e0 | 27 | tcx: TyCtxt<'tcx>, |
60c5eb7d XL |
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) { | |
74b04a01 | 34 | let stmt_before_term = (!block_data.statements.is_empty()) |
60c5eb7d XL |
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) { | |
136023e0 XL |
40 | let ty = place.ty(body, tcx).ty; |
41 | if ty.is_enum() { | |
42 | return Some(ty); | |
60c5eb7d XL |
43 | } |
44 | } | |
45 | } | |
46 | } | |
47 | ||
48 | None | |
49 | } | |
50 | ||
51 | fn variant_discriminants<'tcx>( | |
ba9703b0 | 52 | layout: &TyAndLayout<'tcx>, |
60c5eb7d XL |
53 | ty: Ty<'tcx>, |
54 | tcx: TyCtxt<'tcx>, | |
29967ef6 | 55 | ) -> FxHashSet<u128> { |
ba9703b0 | 56 | match &layout.variants { |
29967ef6 XL |
57 | Variants::Single { index } => { |
58 | let mut res = FxHashSet::default(); | |
5099ac24 FG |
59 | res.insert( |
60 | ty.discriminant_for_variant(tcx, *index) | |
61 | .map_or(index.as_u32() as u128, |discr| discr.val), | |
62 | ); | |
29967ef6 XL |
63 | res |
64 | } | |
60c5eb7d XL |
65 | Variants::Multiple { variants, .. } => variants |
66 | .iter_enumerated() | |
67 | .filter_map(|(idx, layout)| { | |
5e7ed085 | 68 | (layout.abi() != Abi::Uninhabited) |
60c5eb7d XL |
69 | .then(|| ty.discriminant_for_variant(tcx, idx).unwrap().val) |
70 | }) | |
71 | .collect(), | |
72 | } | |
73 | } | |
74 | ||
5e7ed085 FG |
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 | ||
60c5eb7d | 97 | impl<'tcx> MirPass<'tcx> for UninhabitedEnumBranching { |
a2a8927a XL |
98 | fn is_enabled(&self, sess: &rustc_session::Session) -> bool { |
99 | sess.mir_opt_level() > 0 | |
100 | } | |
101 | ||
29967ef6 | 102 | fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
29967ef6 | 103 | trace!("UninhabitedEnumBranching starting for {:?}", body.source); |
60c5eb7d | 104 | |
5099ac24 | 105 | for bb in body.basic_blocks().indices() { |
60c5eb7d XL |
106 | trace!("processing block {:?}", bb); |
107 | ||
3c0e092e XL |
108 | let Some(discriminant_ty) = get_switched_on_type(&body.basic_blocks()[bb], tcx, body) else { |
109 | continue; | |
110 | }; | |
60c5eb7d | 111 | |
29967ef6 | 112 | let layout = tcx.layout_of(tcx.param_env(body.source.def_id()).and(discriminant_ty)); |
60c5eb7d XL |
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 | ||
29967ef6 | 122 | if let TerminatorKind::SwitchInt { targets, .. } = |
60c5eb7d XL |
123 | &mut body.basic_blocks_mut()[bb].terminator_mut().kind |
124 | { | |
5e7ed085 | 125 | let mut new_targets = SwitchTargets::new( |
29967ef6 XL |
126 | targets.iter().filter(|(val, _)| allowed_variants.contains(val)), |
127 | targets.otherwise(), | |
128 | ); | |
129 | ||
5e7ed085 FG |
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 | } | |
60c5eb7d XL |
144 | } else { |
145 | unreachable!() | |
146 | } | |
147 | } | |
148 | } | |
149 | } |