Skip to content

Commit c296b2d

Browse files
borsbors
committed
Auto merge of #65694 - wesleywiser:uninhabited_enum_variants_pass, r=oli-obk
[mir-opt] Implement pass to remove branches on uninhabited variants Based on discussion [here](#64890 (comment)), this is a pass to eliminate dead code that is caused by branching on an enum with uninhabited variants. r? @oli-obk
2 parents ac162c6 + cbe2f60 commit c296b2d

File tree

4 files changed

+354
-0
lines changed

4 files changed

+354
-0
lines changed

‎src/librustc_mir/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ Rust MIR: a lowered representation of Rust. Also: an experiment!
2626
#![feature(associated_type_bounds)]
2727
#![feature(range_is_empty)]
2828
#![feature(stmt_expr_attributes)]
29+
#![feature(bool_to_option)]
2930

3031
#![recursion_limit="256"]
3132

‎src/librustc_mir/transform/mod.rs

+3
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ pub mod const_prop;
3636
pub mod generator;
3737
pub mod inline;
3838
pub mod uniform_array_move_out;
39+
pub mod uninhabited_enum_branching;
3940

4041
pub(crate) fn provide(providers: &mut Providers<'_>) {
4142
self::qualify_consts::provide(providers);
@@ -258,6 +259,8 @@ fn run_optimization_passes<'tcx>(
258259

259260

260261
// Optimizations begin.
262+
&uninhabited_enum_branching::UninhabitedEnumBranching,
263+
&simplify::SimplifyCfg::new("after-uninhabited-enum-branching"),
261264
&uniform_array_move_out::RestoreSubsliceArrayMoveOut::new(tcx),
262265
&inline::Inline,
263266

‎src/librustc_mir/transform/uninhabited_enum_branching.rs

+126
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
//! A pass that eliminates branches on uninhabited enum variants.
2+
3+
use crate::transform::{MirPass, MirSource};
4+
use rustc::mir::{
5+
BasicBlock, BasicBlockData, Body, Local, Operand, Rvalue, StatementKind, TerminatorKind,
6+
};
7+
use rustc::ty::layout::{Abi, TyLayout, Variants};
8+
use rustc::ty::{Ty, TyCtxt};
9+
10+
pub struct UninhabitedEnumBranching;
11+
12+
fn get_discriminant_local(terminator: &TerminatorKind<'_>) -> Option<Local> {
13+
if let TerminatorKind::SwitchInt { discr: Operand::Move(p), .. } = terminator {
14+
p.as_local()
15+
} else {
16+
None
17+
}
18+
}
19+
20+
/// If the basic block terminates by switching on a discriminant, this returns the `Ty` the
21+
/// discriminant is read from. Otherwise, returns None.
22+
fn get_switched_on_type<'tcx>(
23+
block_data: &BasicBlockData<'tcx>,
24+
body: &Body<'tcx>,
25+
) -> Option<Ty<'tcx>> {
26+
let terminator = block_data.terminator();
27+
28+
// Only bother checking blocks which terminate by switching on a local.
29+
if let Some(local) = get_discriminant_local(&terminator.kind) {
30+
let stmt_before_term = (block_data.statements.len() > 0)
31+
.then_with(|| &block_data.statements[block_data.statements.len() - 1].kind);
32+
33+
if let Some(StatementKind::Assign(box (l, Rvalue::Discriminant(place)))) = stmt_before_term
34+
{
35+
if l.as_local() == Some(local) {
36+
if let Some(r_local) = place.as_local() {
37+
let ty = body.local_decls[r_local].ty;
38+
39+
if ty.is_enum() {
40+
return Some(ty);
41+
}
42+
}
43+
}
44+
}
45+
}
46+
47+
None
48+
}
49+
50+
fn variant_discriminants<'tcx>(
51+
layout: &TyLayout<'tcx>,
52+
ty: Ty<'tcx>,
53+
tcx: TyCtxt<'tcx>,
54+
) -> Vec<u128> {
55+
match &layout.details.variants {
56+
Variants::Single { index } => vec![index.as_u32() as u128],
57+
Variants::Multiple { variants, .. } => variants
58+
.iter_enumerated()
59+
.filter_map(|(idx, layout)| {
60+
(layout.abi != Abi::Uninhabited)
61+
.then_with(|| ty.discriminant_for_variant(tcx, idx).unwrap().val)
62+
})
63+
.collect(),
64+
}
65+
}
66+
67+
impl<'tcx> MirPass<'tcx> for UninhabitedEnumBranching {
68+
fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut Body<'tcx>) {
69+
if source.promoted.is_some() {
70+
return;
71+
}
72+
73+
trace!("UninhabitedEnumBranching starting for {:?}", source);
74+
75+
let basic_block_count = body.basic_blocks().len();
76+
77+
for bb in 0..basic_block_count {
78+
let bb = BasicBlock::from_usize(bb);
79+
trace!("processing block {:?}", bb);
80+
81+
let discriminant_ty =
82+
if let Some(ty) = get_switched_on_type(&body.basic_blocks()[bb], body) {
83+
ty
84+
} else {
85+
continue;
86+
};
87+
88+
let layout = tcx.layout_of(tcx.param_env(source.def_id()).and(discriminant_ty));
89+
90+
let allowed_variants = if let Ok(layout) = layout {
91+
variant_discriminants(&layout, discriminant_ty, tcx)
92+
} else {
93+
continue;
94+
};
95+
96+
trace!("allowed_variants = {:?}", allowed_variants);
97+
98+
if let TerminatorKind::SwitchInt { values, targets, .. } =
99+
&mut body.basic_blocks_mut()[bb].terminator_mut().kind
100+
{
101+
let vals = &*values;
102+
let zipped = vals.iter().zip(targets.into_iter());
103+
104+
let mut matched_values = Vec::with_capacity(allowed_variants.len());
105+
let mut matched_targets = Vec::with_capacity(allowed_variants.len() + 1);
106+
107+
for (val, target) in zipped {
108+
if allowed_variants.contains(val) {
109+
matched_values.push(*val);
110+
matched_targets.push(*target);
111+
} else {
112+
trace!("eliminating {:?} -> {:?}", val, target);
113+
}
114+
}
115+
116+
// handle the "otherwise" branch
117+
matched_targets.push(targets.pop().unwrap());
118+
119+
*values = matched_values.into();
120+
*targets = matched_targets;
121+
} else {
122+
unreachable!()
123+
}
124+
}
125+
}
126+
}

‎src/test/mir-opt/uninhabited_enum_branching.rs

+224
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,224 @@
1+
enum Empty { }
2+
3+
// test matching an enum with uninhabited variants
4+
enum Test1 {
5+
A(Empty),
6+
B(Empty),
7+
C
8+
}
9+
10+
// test an enum where the discriminants don't match the variant indexes
11+
// (the optimization should do nothing here)
12+
enum Test2 {
13+
D = 4,
14+
E = 5,
15+
}
16+
17+
fn main() {
18+
match Test1::C {
19+
Test1::A(_) => "A(Empty)",
20+
Test1::B(_) => "B(Empty)",
21+
Test1::C => "C",
22+
};
23+
24+
match Test2::D {
25+
Test2::D => "D",
26+
Test2::E => "E",
27+
};
28+
}
29+
30+
// END RUST SOURCE
31+
//
32+
// START rustc.main.UninhabitedEnumBranching.before.mir
33+
// let mut _0: ();
34+
// let _1: &str;
35+
// let mut _2: Test1;
36+
// let mut _3: isize;
37+
// let mut _4: &str;
38+
// let mut _5: &str;
39+
// let _6: &str;
40+
// let mut _7: Test2;
41+
// let mut _8: isize;
42+
// let mut _9: &str;
43+
// bb0: {
44+
// StorageLive(_1);
45+
// StorageLive(_2);
46+
// _2 = Test1::C;
47+
// _3 = discriminant(_2);
48+
// switchInt(move _3) -> [0isize: bb3, 1isize: bb4, 2isize: bb1, otherwise: bb2];
49+
// }
50+
// bb1: {
51+
// StorageLive(_5);
52+
// _5 = const "C";
53+
// _1 = &(*_5);
54+
// StorageDead(_5);
55+
// goto -> bb5;
56+
// }
57+
// bb2: {
58+
// unreachable;
59+
// }
60+
// bb3: {
61+
// _1 = const "A(Empty)";
62+
// goto -> bb5;
63+
// }
64+
// bb4: {
65+
// StorageLive(_4);
66+
// _4 = const "B(Empty)";
67+
// _1 = &(*_4);
68+
// StorageDead(_4);
69+
// goto -> bb5;
70+
// }
71+
// bb5: {
72+
// StorageDead(_2);
73+
// StorageDead(_1);
74+
// StorageLive(_6);
75+
// StorageLive(_7);
76+
// _7 = Test2::D;
77+
// _8 = discriminant(_7);
78+
// switchInt(move _8) -> [4isize: bb8, 5isize: bb6, otherwise: bb7];
79+
// }
80+
// bb6: {
81+
// StorageLive(_9);
82+
// _9 = const "E";
83+
// _6 = &(*_9);
84+
// StorageDead(_9);
85+
// goto -> bb9;
86+
// }
87+
// bb7: {
88+
// unreachable;
89+
// }
90+
// bb8: {
91+
// _6 = const "D";
92+
// goto -> bb9;
93+
// }
94+
// bb9: {
95+
// StorageDead(_7);
96+
// StorageDead(_6);
97+
// _0 = ();
98+
// return;
99+
// }
100+
// END rustc.main.UninhabitedEnumBranching.before.mir
101+
// START rustc.main.UninhabitedEnumBranching.after.mir
102+
// let mut _0: ();
103+
// let _1: &str;
104+
// let mut _2: Test1;
105+
// let mut _3: isize;
106+
// let mut _4: &str;
107+
// let mut _5: &str;
108+
// let _6: &str;
109+
// let mut _7: Test2;
110+
// let mut _8: isize;
111+
// let mut _9: &str;
112+
// bb0: {
113+
// StorageLive(_1);
114+
// StorageLive(_2);
115+
// _2 = Test1::C;
116+
// _3 = discriminant(_2);
117+
// switchInt(move _3) -> [2isize: bb1, otherwise: bb2];
118+
// }
119+
// bb1: {
120+
// StorageLive(_5);
121+
// _5 = const "C";
122+
// _1 = &(*_5);
123+
// StorageDead(_5);
124+
// goto -> bb5;
125+
// }
126+
// bb2: {
127+
// unreachable;
128+
// }
129+
// bb3: {
130+
// _1 = const "A(Empty)";
131+
// goto -> bb5;
132+
// }
133+
// bb4: {
134+
// StorageLive(_4);
135+
// _4 = const "B(Empty)";
136+
// _1 = &(*_4);
137+
// StorageDead(_4);
138+
// goto -> bb5;
139+
// }
140+
// bb5: {
141+
// StorageDead(_2);
142+
// StorageDead(_1);
143+
// StorageLive(_6);
144+
// StorageLive(_7);
145+
// _7 = Test2::D;
146+
// _8 = discriminant(_7);
147+
// switchInt(move _8) -> [4isize: bb8, 5isize: bb6, otherwise: bb7];
148+
// }
149+
// bb6: {
150+
// StorageLive(_9);
151+
// _9 = const "E";
152+
// _6 = &(*_9);
153+
// StorageDead(_9);
154+
// goto -> bb9;
155+
// }
156+
// bb7: {
157+
// unreachable;
158+
// }
159+
// bb8: {
160+
// _6 = const "D";
161+
// goto -> bb9;
162+
// }
163+
// bb9: {
164+
// StorageDead(_7);
165+
// StorageDead(_6);
166+
// _0 = ();
167+
// return;
168+
// }
169+
// END rustc.main.UninhabitedEnumBranching.after.mir
170+
// START rustc.main.SimplifyCfg-after-uninhabited-enum-branching.after.mir
171+
// let mut _0: ();
172+
// let _1: &str;
173+
// let mut _2: Test1;
174+
// let mut _3: isize;
175+
// let mut _4: &str;
176+
// let mut _5: &str;
177+
// let _6: &str;
178+
// let mut _7: Test2;
179+
// let mut _8: isize;
180+
// let mut _9: &str;
181+
// bb0: {
182+
// StorageLive(_1);
183+
// StorageLive(_2);
184+
// _2 = Test1::C;
185+
// _3 = discriminant(_2);
186+
// switchInt(move _3) -> [2isize: bb1, otherwise: bb2];
187+
// }
188+
// bb1: {
189+
// StorageLive(_5);
190+
// _5 = const "C";
191+
// _1 = &(*_5);
192+
// StorageDead(_5);
193+
// StorageDead(_2);
194+
// StorageDead(_1);
195+
// StorageLive(_6);
196+
// StorageLive(_7);
197+
// _7 = Test2::D;
198+
// _8 = discriminant(_7);
199+
// switchInt(move _8) -> [4isize: bb5, 5isize: bb3, otherwise: bb4];
200+
// }
201+
// bb2: {
202+
// unreachable;
203+
// }
204+
// bb3: {
205+
// StorageLive(_9);
206+
// _9 = const "E";
207+
// _6 = &(*_9);
208+
// StorageDead(_9);
209+
// goto -> bb6;
210+
// }
211+
// bb4: {
212+
// unreachable;
213+
// }
214+
// bb5: {
215+
// _6 = const "D";
216+
// goto -> bb6;
217+
// }
218+
// bb6: {
219+
// StorageDead(_7);
220+
// StorageDead(_6);
221+
// _0 = ();
222+
// return;
223+
// }
224+
// END rustc.main.SimplifyCfg-after-uninhabited-enum-branching.after.mir

0 commit comments

Comments
 (0)