Skip to content

Commit f9395e1

Browse files
cjgillotcjgillot
committed
Replace SwitchInt to unreachable by an assumption.
1 parent 459cfcc commit f9395e1

11 files changed

+200
-114
lines changed

‎compiler/rustc_mir_transform/src/unreachable_prop.rs

+117-82
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,12 @@
33
//! post-order traversal of the blocks.
44
55
use crate::MirPass;
6-
use rustc_data_structures::fx::{FxHashMap, FxHashSet};
6+
use rustc_data_structures::fx::FxHashSet;
7+
use rustc_middle::mir::interpret::Scalar;
8+
use rustc_middle::mir::patch::MirPatch;
79
use rustc_middle::mir::*;
8-
use rustc_middle::ty::TyCtxt;
10+
use rustc_middle::ty::{self, TyCtxt};
11+
use rustc_target::abi::Size;
912

1013
pub struct UnreachablePropagation;
1114

@@ -16,102 +19,134 @@ impl MirPass<'_> for UnreachablePropagation {
1619
}
1720

1821
fn run_pass<'tcx>(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
22+
let mut patch = MirPatch::new(body);
1923
let mut unreachable_blocks = FxHashSet::default();
20-
let mut replacements = FxHashMap::default();
2124

2225
for (bb, bb_data) in traversal::postorder(body) {
2326
let terminator = bb_data.terminator();
24-
if terminator.kind == TerminatorKind::Unreachable {
25-
unreachable_blocks.insert(bb);
26-
} else {
27-
let is_unreachable = |succ: BasicBlock| unreachable_blocks.contains(&succ);
28-
let terminator_kind_opt = remove_successors(&terminator.kind, is_unreachable);
29-
30-
if let Some(terminator_kind) = terminator_kind_opt {
31-
if terminator_kind == TerminatorKind::Unreachable {
32-
unreachable_blocks.insert(bb);
33-
}
34-
replacements.insert(bb, terminator_kind);
27+
let is_unreachable = match &terminator.kind {
28+
TerminatorKind::Unreachable => true,
29+
// This will unconditionally run into an unreachable and is therefore unreachable as well.
30+
TerminatorKind::Goto { target } if unreachable_blocks.contains(target) => {
31+
patch.patch_terminator(bb, TerminatorKind::Unreachable);
32+
true
33+
}
34+
// Try to remove unreachable targets from the switch.
35+
TerminatorKind::SwitchInt { .. } => {
36+
remove_successors_from_switch(tcx, bb, &unreachable_blocks, body, &mut patch)
3537
}
38+
_ => false,
39+
};
40+
if is_unreachable {
41+
unreachable_blocks.insert(bb);
3642
}
3743
}
3844

39-
// We do want do keep some unreachable blocks, but make them empty.
40-
for bb in unreachable_blocks {
41-
if !tcx.consider_optimizing(|| {
42-
format!("UnreachablePropagation {:?} ", body.source.def_id())
43-
}) {
44-
break;
45-
}
46-
47-
body.basic_blocks_mut()[bb].statements.clear();
45+
if !tcx
46+
.consider_optimizing(|| format!("UnreachablePropagation {:?} ", body.source.def_id()))
47+
{
48+
return;
4849
}
4950

50-
for (bb, terminator_kind) in replacements {
51-
if !tcx.consider_optimizing(|| {
52-
format!("UnreachablePropagation {:?} ", body.source.def_id())
53-
}) {
54-
break;
55-
}
51+
patch.apply(body);
5652

57-
body.basic_blocks_mut()[bb].terminator_mut().kind = terminator_kind;
53+
// We do want do keep some unreachable blocks, but make them empty.
54+
for bb in unreachable_blocks {
55+
body.basic_blocks_mut()[bb].statements.clear();
5856
}
59-
60-
// Do not remove dead blocks, let `SimplifyCfg` do it.
6157
}
6258
}
6359

64-
fn remove_successors<'tcx, F>(
65-
terminator_kind: &TerminatorKind<'tcx>,
66-
is_unreachable: F,
67-
) -> Option<TerminatorKind<'tcx>>
68-
where
69-
F: Fn(BasicBlock) -> bool,
70-
{
71-
let terminator = match terminator_kind {
72-
// This will unconditionally run into an unreachable and is therefore unreachable as well.
73-
TerminatorKind::Goto { target } if is_unreachable(*target) => TerminatorKind::Unreachable,
74-
TerminatorKind::SwitchInt { targets, discr } => {
75-
let otherwise = targets.otherwise();
76-
77-
// If all targets are unreachable, we can be unreachable as well.
78-
if targets.all_targets().iter().all(|bb| is_unreachable(*bb)) {
79-
TerminatorKind::Unreachable
80-
} else if is_unreachable(otherwise) {
81-
// If there are multiple targets, don't delete unreachable branches (like an unreachable otherwise)
82-
// unless otherwise is unreachable, in which case deleting a normal branch causes it to be merged with
83-
// the otherwise, keeping its unreachable.
84-
// This looses information about reachability causing worse codegen.
85-
// For example (see tests/codegen/match-optimizes-away.rs)
86-
//
87-
// pub enum Two { A, B }
88-
// pub fn identity(x: Two) -> Two {
89-
// match x {
90-
// Two::A => Two::A,
91-
// Two::B => Two::B,
92-
// }
93-
// }
94-
//
95-
// This generates a `switchInt() -> [0: 0, 1: 1, otherwise: unreachable]`, which allows us or LLVM to
96-
// turn it into just `x` later. Without the unreachable, such a transformation would be illegal.
97-
// If the otherwise branch is unreachable, we can delete all other unreachable targets, as they will
98-
// still point to the unreachable and therefore not lose reachability information.
99-
let reachable_iter = targets.iter().filter(|(_, bb)| !is_unreachable(*bb));
100-
101-
let new_targets = SwitchTargets::new(reachable_iter, otherwise);
102-
103-
// No unreachable branches were removed.
104-
if new_targets.all_targets().len() == targets.all_targets().len() {
105-
return None;
106-
}
60+
/// Return whether the current terminator is fully unreachable.
61+
fn remove_successors_from_switch<'tcx>(
62+
tcx: TyCtxt<'tcx>,
63+
bb: BasicBlock,
64+
unreachable_blocks: &FxHashSet<BasicBlock>,
65+
body: &Body<'tcx>,
66+
patch: &mut MirPatch<'tcx>,
67+
) -> bool {
68+
let terminator = body.basic_blocks[bb].terminator();
69+
let TerminatorKind::SwitchInt { discr, targets } = &terminator.kind else { bug!() };
70+
let source_info = terminator.source_info;
71+
let location = body.terminator_loc(bb);
72+
73+
let is_unreachable = |bb| unreachable_blocks.contains(&bb);
74+
75+
// If there are multiple targets, we want to keep information about reachability for codegen.
76+
// For example (see tests/codegen/match-optimizes-away.rs)
77+
//
78+
// pub enum Two { A, B }
79+
// pub fn identity(x: Two) -> Two {
80+
// match x {
81+
// Two::A => Two::A,
82+
// Two::B => Two::B,
83+
// }
84+
// }
85+
//
86+
// This generates a `switchInt() -> [0: 0, 1: 1, otherwise: unreachable]`, which allows us or LLVM to
87+
// turn it into just `x` later. Without the unreachable, such a transformation would be illegal.
88+
//
89+
// In order to preserve this information, we record reachable and unreachable targets as
90+
// `Assume` statements in MIR.
91+
92+
let discr_ty = discr.ty(body, tcx);
93+
let discr_size = Size::from_bits(match discr_ty.kind() {
94+
ty::Uint(uint) => uint.normalize(tcx.sess.target.pointer_width).bit_width().unwrap(),
95+
ty::Int(int) => int.normalize(tcx.sess.target.pointer_width).bit_width().unwrap(),
96+
ty::Char => 32,
97+
ty::Bool => 1,
98+
other => bug!("unhandled type: {:?}", other),
99+
});
100+
101+
let mut add_assumption = |binop, value| {
102+
let local = patch.new_temp(tcx.types.bool, source_info.span);
103+
let value = Operand::Constant(Box::new(ConstOperand {
104+
span: source_info.span,
105+
user_ty: None,
106+
const_: Const::from_scalar(tcx, Scalar::from_uint(value, discr_size), discr_ty),
107+
}));
108+
let cmp = Rvalue::BinaryOp(binop, Box::new((discr.to_copy(), value)));
109+
patch.add_assign(location, local.into(), cmp);
110+
111+
let assume = NonDivergingIntrinsic::Assume(Operand::Move(local.into()));
112+
patch.add_statement(location, StatementKind::Intrinsic(Box::new(assume)));
113+
};
107114

108-
TerminatorKind::SwitchInt { discr: discr.clone(), targets: new_targets }
109-
} else {
110-
// If the otherwise branch is reachable, we don't want to delete any unreachable branches.
111-
return None;
112-
}
115+
let reachable_iter = targets.iter().filter(|&(value, bb)| {
116+
let is_unreachable = is_unreachable(bb);
117+
if is_unreachable {
118+
// We remove this target from the switch, so record the inequality using `Assume`.
119+
add_assumption(BinOp::Ne, value);
120+
false
121+
} else {
122+
true
123+
}
124+
});
125+
126+
let otherwise = targets.otherwise();
127+
let new_targets = SwitchTargets::new(reachable_iter, otherwise);
128+
129+
let num_targets = new_targets.all_targets().len();
130+
let otherwise_unreachable = is_unreachable(otherwise);
131+
let fully_unreachable = num_targets == 1 && otherwise_unreachable;
132+
133+
let terminator = match (num_targets, otherwise_unreachable) {
134+
// If all targets are unreachable, we can be unreachable as well.
135+
(1, true) => TerminatorKind::Unreachable,
136+
(1, false) => TerminatorKind::Goto { target: otherwise },
137+
(2, true) => {
138+
// All targets are unreachable except one. Record the equality, and make it a goto.
139+
let (value, target) = new_targets.iter().next().unwrap();
140+
add_assumption(BinOp::Eq, value);
141+
TerminatorKind::Goto { target }
113142
}
114-
_ => return None,
143+
_ if num_targets == targets.all_targets().len() => {
144+
// Nothing has changed.
145+
return false;
146+
}
147+
_ => TerminatorKind::SwitchInt { discr: discr.clone(), targets: new_targets },
115148
};
116-
Some(terminator)
149+
150+
patch.patch_terminator(bb, terminator);
151+
fully_unreachable
117152
}

‎tests/mir-opt/inline/unwrap_unchecked.unwrap_unchecked.Inline.panic-abort.diff

+10-2
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
+ scope 5 (inlined unreachable_unchecked) {
1717
+ scope 6 {
1818
+ scope 7 (inlined unreachable_unchecked::runtime) {
19+
+ let _5: !;
1920
+ }
2021
+ }
2122
+ }
@@ -30,16 +31,23 @@
3031
_2 = move _1;
3132
- _0 = Option::<T>::unwrap_unchecked(move _2) -> [return: bb1, unwind unreachable];
3233
+ StorageLive(_3);
34+
+ StorageLive(_5);
3335
+ _4 = discriminant(_2);
34-
+ switchInt(move _4) -> [1: bb2, otherwise: bb1];
36+
+ switchInt(move _4) -> [0: bb1, 1: bb3, otherwise: bb2];
3537
}
3638

3739
bb1: {
38-
+ unreachable;
40+
+ assume(const false);
41+
+ _5 = core::panicking::panic_nounwind(const "unsafe precondition(s) violated: hint::unreachable_unchecked must never be reached") -> unwind unreachable;
3942
+ }
4043
+
4144
+ bb2: {
45+
+ unreachable;
46+
+ }
47+
+
48+
+ bb3: {
4249
+ _0 = move ((_2 as Some).0: T);
50+
+ StorageDead(_5);
4351
+ StorageDead(_3);
4452
StorageDead(_2);
4553
return;

‎tests/mir-opt/inline/unwrap_unchecked.unwrap_unchecked.Inline.panic-unwind.diff

+10-2
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
+ scope 5 (inlined unreachable_unchecked) {
1717
+ scope 6 {
1818
+ scope 7 (inlined unreachable_unchecked::runtime) {
19+
+ let _5: !;
1920
+ }
2021
+ }
2122
+ }
@@ -30,20 +31,27 @@
3031
_2 = move _1;
3132
- _0 = Option::<T>::unwrap_unchecked(move _2) -> [return: bb1, unwind: bb2];
3233
+ StorageLive(_3);
34+
+ StorageLive(_5);
3335
+ _4 = discriminant(_2);
34-
+ switchInt(move _4) -> [1: bb2, otherwise: bb1];
36+
+ switchInt(move _4) -> [0: bb1, 1: bb3, otherwise: bb2];
3537
}
3638

3739
bb1: {
3840
- StorageDead(_2);
3941
- return;
40-
+ unreachable;
42+
+ assume(const false);
43+
+ _5 = core::panicking::panic_nounwind(const "unsafe precondition(s) violated: hint::unreachable_unchecked must never be reached") -> unwind unreachable;
4144
}
4245

4346
- bb2 (cleanup): {
4447
- resume;
4548
+ bb2: {
49+
+ unreachable;
50+
+ }
51+
+
52+
+ bb3: {
4653
+ _0 = move ((_2 as Some).0: T);
54+
+ StorageDead(_5);
4755
+ StorageDead(_3);
4856
+ StorageDead(_2);
4957
+ return;

‎tests/mir-opt/inline/unwrap_unchecked.unwrap_unchecked.PreCodegen.after.panic-abort.mir

+12-6
Original file line numberDiff line numberDiff line change
@@ -6,36 +6,42 @@ fn unwrap_unchecked(_1: Option<T>) -> T {
66
scope 1 (inlined #[track_caller] Option::<T>::unwrap_unchecked) {
77
debug self => _1;
88
let mut _2: isize;
9-
let mut _3: &std::option::Option<T>;
9+
let mut _4: &std::option::Option<T>;
1010
scope 2 {
1111
debug val => _0;
1212
}
1313
scope 3 {
1414
scope 5 (inlined unreachable_unchecked) {
1515
scope 6 {
1616
scope 7 (inlined unreachable_unchecked::runtime) {
17+
let _3: !;
1718
}
1819
}
1920
}
2021
}
2122
scope 4 (inlined Option::<T>::is_some) {
22-
debug self => _3;
23+
debug self => _4;
2324
}
2425
}
2526

2627
bb0: {
27-
StorageLive(_3);
28+
StorageLive(_4);
2829
_2 = discriminant(_1);
29-
switchInt(move _2) -> [1: bb1, otherwise: bb2];
30+
switchInt(move _2) -> [0: bb1, 1: bb2, otherwise: bb3];
3031
}
3132

3233
bb1: {
34+
assume(const false);
35+
_3 = core::panicking::panic_nounwind(const "unsafe precondition(s) violated: hint::unreachable_unchecked must never be reached") -> unwind unreachable;
36+
}
37+
38+
bb2: {
3339
_0 = move ((_1 as Some).0: T);
34-
StorageDead(_3);
40+
StorageDead(_4);
3541
return;
3642
}
3743

38-
bb2: {
44+
bb3: {
3945
unreachable;
4046
}
4147
}

‎tests/mir-opt/inline/unwrap_unchecked.unwrap_unchecked.PreCodegen.after.panic-unwind.mir

+12-6
Original file line numberDiff line numberDiff line change
@@ -6,36 +6,42 @@ fn unwrap_unchecked(_1: Option<T>) -> T {
66
scope 1 (inlined #[track_caller] Option::<T>::unwrap_unchecked) {
77
debug self => _1;
88
let mut _2: isize;
9-
let mut _3: &std::option::Option<T>;
9+
let mut _4: &std::option::Option<T>;
1010
scope 2 {
1111
debug val => _0;
1212
}
1313
scope 3 {
1414
scope 5 (inlined unreachable_unchecked) {
1515
scope 6 {
1616
scope 7 (inlined unreachable_unchecked::runtime) {
17+
let _3: !;
1718
}
1819
}
1920
}
2021
}
2122
scope 4 (inlined Option::<T>::is_some) {
22-
debug self => _3;
23+
debug self => _4;
2324
}
2425
}
2526

2627
bb0: {
27-
StorageLive(_3);
28+
StorageLive(_4);
2829
_2 = discriminant(_1);
29-
switchInt(move _2) -> [1: bb1, otherwise: bb2];
30+
switchInt(move _2) -> [0: bb1, 1: bb2, otherwise: bb3];
3031
}
3132

3233
bb1: {
34+
assume(const false);
35+
_3 = core::panicking::panic_nounwind(const "unsafe precondition(s) violated: hint::unreachable_unchecked must never be reached") -> unwind unreachable;
36+
}
37+
38+
bb2: {
3339
_0 = move ((_1 as Some).0: T);
34-
StorageDead(_3);
40+
StorageDead(_4);
3541
return;
3642
}
3743

38-
bb2: {
44+
bb3: {
3945
unreachable;
4046
}
4147
}

0 commit comments

Comments
 (0)