Skip to content

Commit a77e2ad

Browse files
MoxinilianMoxinilian
committed
implemented separate_const_switch MIR opt
un-update itertools improve predecessor amount short-circuiting cleanup and comments somewhat improved drawing
1 parent b8be316 commit a77e2ad

9 files changed

+1055
-0
lines changed

‎compiler/rustc_mir/src/transform/mod.rs

+2
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ pub mod remove_unneeded_drops;
4848
pub mod remove_zsts;
4949
pub mod required_consts;
5050
pub mod rustc_peek;
51+
pub mod separate_const_switch;
5152
pub mod simplify;
5253
pub mod simplify_branches;
5354
pub mod simplify_comparison_integral;
@@ -501,6 +502,7 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
501502
// inst combine is after MatchBranchSimplification to clean up Ne(_1, false)
502503
&multiple_return_terminators::MultipleReturnTerminators,
503504
&instcombine::InstCombine,
505+
&separate_const_switch::SeparateConstSwitch,
504506
&const_prop::ConstProp,
505507
&simplify_branches::SimplifyBranches::new("after-const-prop"),
506508
&early_otherwise_branch::EarlyOtherwiseBranch,

‎compiler/rustc_mir/src/transform/separate_const_switch.rs

+343
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,343 @@
1+
//! A pass that duplicates switch-terminated blocks
2+
//! into a new copy for each predecessor, provided
3+
//! the predecessor sets the value being switched
4+
//! over to a constant.
5+
//!
6+
//! The purpose of this pass is to help constant
7+
//! propagation passes to simplify the switch terminator
8+
//! of the copied blocks into gotos when some predecessors
9+
//! statically determine the output of switches.
10+
//!
11+
//! ```text
12+
//! x = 12 --- ---> something
13+
//! \ / 12
14+
//! --> switch x
15+
//! / \ otherwise
16+
//! x = y --- ---> something else
17+
//! ```
18+
//! becomes
19+
//! ```text
20+
//! x = 12 ---> switch x ------> something
21+
//! \ / 12
22+
//! X
23+
//! / \ otherwise
24+
//! x = y ---> switch x ------> something else
25+
//! ```
26+
//! so it can hopefully later be turned by another pass into
27+
//! ```text
28+
//! x = 12 --------------------> something
29+
//! / 12
30+
//! /
31+
//! / otherwise
32+
//! x = y ---- switch x ------> something else
33+
//! ```
34+
//!
35+
//! This optimization is meant to cover simple cases
36+
//! like `?` desugaring. For now, it thus focuses on
37+
//! simplicity rather than completeness (it notably
38+
//! sometimes duplicates abusively).
39+
40+
use crate::transform::MirPass;
41+
use rustc_middle::mir::*;
42+
use rustc_middle::ty::TyCtxt;
43+
use smallvec::SmallVec;
44+
45+
pub struct SeparateConstSwitch;
46+
47+
impl<'tcx> MirPass<'tcx> for SeparateConstSwitch {
48+
fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
49+
if tcx.sess.mir_opt_level() < 4 {
50+
return;
51+
}
52+
53+
// If execution did something, applying a simplification layer
54+
// helps later passes optimize the copy away.
55+
if separate_const_switch(body) > 0 {
56+
super::simplify::simplify_cfg(tcx, body);
57+
}
58+
}
59+
}
60+
61+
/// Returns the amount of blocks that were duplicated
62+
pub fn separate_const_switch<'tcx>(body: &mut Body<'tcx>) -> usize {
63+
let mut new_blocks: SmallVec<[(BasicBlock, BasicBlock); 6]> = SmallVec::new();
64+
let predecessors = body.predecessors();
65+
'block_iter: for (block_id, block) in body.basic_blocks().iter_enumerated() {
66+
if let TerminatorKind::SwitchInt {
67+
discr: Operand::Copy(switch_place) | Operand::Move(switch_place),
68+
..
69+
} = block.terminator().kind
70+
{
71+
// If the block is on an unwind path, do not
72+
// apply the optimization as unwind paths
73+
// rely on a unique parent invariant
74+
if block.is_cleanup {
75+
continue 'block_iter;
76+
}
77+
78+
// If the block has fewer than 2 predecessors, ignore it
79+
// we could maybe chain blocks that have exactly one
80+
// predecessor, but for now we ignore that
81+
if predecessors[block_id].len() < 2 {
82+
continue 'block_iter;
83+
}
84+
85+
// First, let's find a non-const place
86+
// that determines the result of the switch
87+
if let Some(switch_place) = find_determining_place(switch_place, block) {
88+
// We now have an input place for which it would
89+
// be interesting if predecessors assigned it from a const
90+
91+
let mut predecessors_left = predecessors[block_id].len();
92+
'predec_iter: for predecessor_id in predecessors[block_id].iter().copied() {
93+
let predecessor = &body.basic_blocks()[predecessor_id];
94+
95+
// First we make sure the predecessor jumps
96+
// in a reasonable way
97+
match &predecessor.terminator().kind {
98+
// The following terminators are
99+
// unconditionally valid
100+
TerminatorKind::Goto { .. } | TerminatorKind::SwitchInt { .. } => {}
101+
102+
TerminatorKind::FalseEdge { real_target, .. } => {
103+
if *real_target != block_id {
104+
continue 'predec_iter;
105+
}
106+
}
107+
108+
// The following terminators are not allowed
109+
TerminatorKind::Resume
110+
| TerminatorKind::Drop { .. }
111+
| TerminatorKind::DropAndReplace { .. }
112+
| TerminatorKind::Call { .. }
113+
| TerminatorKind::Assert { .. }
114+
| TerminatorKind::FalseUnwind { .. }
115+
| TerminatorKind::Yield { .. }
116+
| TerminatorKind::Abort
117+
| TerminatorKind::Return
118+
| TerminatorKind::Unreachable
119+
| TerminatorKind::InlineAsm { .. }
120+
| TerminatorKind::GeneratorDrop => {
121+
continue 'predec_iter;
122+
}
123+
}
124+
125+
if is_likely_const(switch_place, predecessor) {
126+
new_blocks.push((predecessor_id, block_id));
127+
predecessors_left -= 1;
128+
if predecessors_left < 2 {
129+
// If the original block only has one predecessor left,
130+
// we have nothing left to do
131+
break 'predec_iter;
132+
}
133+
}
134+
}
135+
}
136+
}
137+
}
138+
139+
// Once the analysis is done, perform the duplication
140+
let body_span = body.span;
141+
let copied_blocks = new_blocks.len();
142+
let blocks = body.basic_blocks_mut();
143+
for (pred_id, target_id) in new_blocks {
144+
let new_block = blocks[target_id].clone();
145+
let new_block_id = blocks.push(new_block);
146+
let terminator = blocks[pred_id].terminator_mut();
147+
148+
match terminator.kind {
149+
TerminatorKind::Goto { ref mut target } => {
150+
*target = new_block_id;
151+
}
152+
153+
TerminatorKind::FalseEdge { ref mut real_target, .. } => {
154+
if *real_target == target_id {
155+
*real_target = new_block_id;
156+
}
157+
}
158+
159+
TerminatorKind::SwitchInt { ref mut targets, .. } => {
160+
targets.all_targets_mut().iter_mut().for_each(|x| {
161+
if *x == target_id {
162+
*x = new_block_id;
163+
}
164+
});
165+
}
166+
167+
TerminatorKind::Resume
168+
| TerminatorKind::Abort
169+
| TerminatorKind::Return
170+
| TerminatorKind::Unreachable
171+
| TerminatorKind::GeneratorDrop
172+
| TerminatorKind::Assert { .. }
173+
| TerminatorKind::DropAndReplace { .. }
174+
| TerminatorKind::FalseUnwind { .. }
175+
| TerminatorKind::Drop { .. }
176+
| TerminatorKind::Call { .. }
177+
| TerminatorKind::InlineAsm { .. }
178+
| TerminatorKind::Yield { .. } => {
179+
span_bug!(
180+
body_span,
181+
"basic block terminator had unexpected kind {:?}",
182+
&terminator.kind
183+
)
184+
}
185+
}
186+
}
187+
188+
copied_blocks
189+
}
190+
191+
/// This function describes a rough heuristic guessing
192+
/// whether a place is last set with a const within the block.
193+
/// Notably, it will be overly pessimistic in cases that are already
194+
/// not handled by `separate_const_switch`.
195+
fn is_likely_const<'tcx>(mut tracked_place: Place<'tcx>, block: &BasicBlockData<'tcx>) -> bool {
196+
for statement in block.statements.iter().rev() {
197+
match &statement.kind {
198+
StatementKind::Assign(assign) => {
199+
if assign.0 == tracked_place {
200+
match assign.1 {
201+
// These rvalues are definitely constant
202+
Rvalue::Use(Operand::Constant(_))
203+
| Rvalue::Ref(_, _, _)
204+
| Rvalue::AddressOf(_, _)
205+
| Rvalue::Cast(_, Operand::Constant(_), _)
206+
| Rvalue::NullaryOp(_, _)
207+
| Rvalue::UnaryOp(_, Operand::Constant(_)) => return true,
208+
209+
// These rvalues make things ambiguous
210+
Rvalue::Repeat(_, _)
211+
| Rvalue::ThreadLocalRef(_)
212+
| Rvalue::Len(_)
213+
| Rvalue::BinaryOp(_, _)
214+
| Rvalue::CheckedBinaryOp(_, _)
215+
| Rvalue::Aggregate(_, _) => return false,
216+
217+
// These rvalues move the place to track
218+
Rvalue::Cast(_, Operand::Copy(place) | Operand::Move(place), _)
219+
| Rvalue::Use(Operand::Copy(place) | Operand::Move(place))
220+
| Rvalue::UnaryOp(_, Operand::Copy(place) | Operand::Move(place))
221+
| Rvalue::Discriminant(place) => tracked_place = place,
222+
}
223+
}
224+
}
225+
226+
// If the discriminant is set, it is always set
227+
// as a constant, so the job is done.
228+
// As we are **ignoring projections**, if the place
229+
// we are tracking sees its discriminant be set,
230+
// that means we had to be tracking the discriminant
231+
// specifically (as it is impossible to switch over
232+
// an enum directly, and if we were switching over
233+
// its content, we would have had to at least cast it to
234+
// some variant first)
235+
StatementKind::SetDiscriminant { place, .. } => {
236+
if **place == tracked_place {
237+
return true;
238+
}
239+
}
240+
241+
// If inline assembly is found, we probably should
242+
// not try to analyze the code
243+
StatementKind::LlvmInlineAsm(_) => return false,
244+
245+
// These statements have no influence on the place
246+
// we are interested in
247+
StatementKind::FakeRead(_)
248+
| StatementKind::StorageLive(_)
249+
| StatementKind::Retag(_, _)
250+
| StatementKind::AscribeUserType(_, _)
251+
| StatementKind::Coverage(_)
252+
| StatementKind::StorageDead(_)
253+
| StatementKind::CopyNonOverlapping(_)
254+
| StatementKind::Nop => {}
255+
}
256+
}
257+
258+
// If no good reason for the place to be const is found,
259+
// give up. We could maybe go up predecessors, but in
260+
// most cases giving up now should be sufficient.
261+
false
262+
}
263+
264+
/// Finds a unique place that entirely determines the value
265+
/// of `switch_place`, if it exists. This is only a heuristic.
266+
/// Ideally we would like to track multiple determining places
267+
/// for some edge cases, but one is enough for a lot of situations.
268+
fn find_determining_place<'tcx>(
269+
mut switch_place: Place<'tcx>,
270+
block: &BasicBlockData<'tcx>,
271+
) -> Option<Place<'tcx>> {
272+
for statement in block.statements.iter().rev() {
273+
match &statement.kind {
274+
StatementKind::Assign(op) => {
275+
if op.0 != switch_place {
276+
continue;
277+
}
278+
279+
match op.1 {
280+
// The following rvalues move the place
281+
// that may be const in the predecessor
282+
Rvalue::Use(Operand::Move(new) | Operand::Copy(new))
283+
| Rvalue::UnaryOp(_, Operand::Copy(new) | Operand::Move(new))
284+
| Rvalue::Cast(_, Operand::Move(new) | Operand::Copy(new), _)
285+
| Rvalue::Repeat(Operand::Move(new) | Operand::Copy(new), _)
286+
| Rvalue::Discriminant(new)
287+
=> switch_place = new,
288+
289+
// The following rvalues might still make the block
290+
// be valid but for now we reject them
291+
Rvalue::Len(_)
292+
| Rvalue::Ref(_, _, _)
293+
| Rvalue::BinaryOp(_, _)
294+
| Rvalue::CheckedBinaryOp(_, _)
295+
| Rvalue::Aggregate(_, _)
296+
297+
// The following rvalues definitely mean we cannot
298+
// or should not apply this optimization
299+
| Rvalue::Use(Operand::Constant(_))
300+
| Rvalue::Repeat(Operand::Constant(_), _)
301+
| Rvalue::ThreadLocalRef(_)
302+
| Rvalue::AddressOf(_, _)
303+
| Rvalue::NullaryOp(_, _)
304+
| Rvalue::UnaryOp(_, Operand::Constant(_))
305+
| Rvalue::Cast(_, Operand::Constant(_), _)
306+
=> return None,
307+
}
308+
}
309+
310+
// These statements have no influence on the place
311+
// we are interested in
312+
StatementKind::FakeRead(_)
313+
| StatementKind::StorageLive(_)
314+
| StatementKind::StorageDead(_)
315+
| StatementKind::Retag(_, _)
316+
| StatementKind::AscribeUserType(_, _)
317+
| StatementKind::Coverage(_)
318+
| StatementKind::CopyNonOverlapping(_)
319+
| StatementKind::Nop => {}
320+
321+
// If inline assembly is found, we probably should
322+
// not try to analyze the code
323+
StatementKind::LlvmInlineAsm(_) => return None,
324+
325+
// If the discriminant is set, it is always set
326+
// as a constant, so the job is already done.
327+
// As we are **ignoring projections**, if the place
328+
// we are tracking sees its discriminant be set,
329+
// that means we had to be tracking the discriminant
330+
// specifically (as it is impossible to switch over
331+
// an enum directly, and if we were switching over
332+
// its content, we would have had to at least cast it to
333+
// some variant first)
334+
StatementKind::SetDiscriminant { place, .. } => {
335+
if **place == switch_place {
336+
return None;
337+
}
338+
}
339+
}
340+
}
341+
342+
Some(switch_place)
343+
}

0 commit comments

Comments
 (0)