]>
Commit | Line | Data |
---|---|---|
6a06907d XL |
1 | //! This pass optimizes the following sequence |
2 | //! ```rust,ignore (example) | |
3 | //! bb2: { | |
4 | //! _2 = const true; | |
5 | //! goto -> bb3; | |
6 | //! } | |
7 | //! | |
8 | //! bb3: { | |
9 | //! switchInt(_2) -> [false: bb4, otherwise: bb5]; | |
10 | //! } | |
11 | //! ``` | |
12 | //! into | |
13 | //! ```rust,ignore (example) | |
14 | //! bb2: { | |
15 | //! _2 = const true; | |
16 | //! goto -> bb5; | |
17 | //! } | |
18 | //! ``` | |
19 | ||
20 | use crate::transform::MirPass; | |
21 | use rustc_middle::mir::*; | |
22 | use rustc_middle::ty::TyCtxt; | |
23 | use rustc_middle::{mir::visit::Visitor, ty::ParamEnv}; | |
24 | ||
25 | use super::simplify::{simplify_cfg, simplify_locals}; | |
26 | ||
27 | pub struct ConstGoto; | |
28 | ||
29 | impl<'tcx> MirPass<'tcx> for ConstGoto { | |
30 | fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { | |
31 | if tcx.sess.mir_opt_level() < 4 { | |
32 | return; | |
33 | } | |
34 | trace!("Running ConstGoto on {:?}", body.source); | |
35 | let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id()); | |
36 | let mut opt_finder = | |
37 | ConstGotoOptimizationFinder { tcx, body, optimizations: vec![], param_env }; | |
38 | opt_finder.visit_body(body); | |
39 | let should_simplify = !opt_finder.optimizations.is_empty(); | |
40 | for opt in opt_finder.optimizations { | |
41 | let terminator = body.basic_blocks_mut()[opt.bb_with_goto].terminator_mut(); | |
42 | let new_goto = TerminatorKind::Goto { target: opt.target_to_use_in_goto }; | |
43 | debug!("SUCCESS: replacing `{:?}` with `{:?}`", terminator.kind, new_goto); | |
44 | terminator.kind = new_goto; | |
45 | } | |
46 | ||
47 | // if we applied optimizations, we potentially have some cfg to cleanup to | |
48 | // make it easier for further passes | |
49 | if should_simplify { | |
50 | simplify_cfg(body); | |
51 | simplify_locals(body, tcx); | |
52 | } | |
53 | } | |
54 | } | |
55 | ||
56 | impl<'a, 'tcx> Visitor<'tcx> for ConstGotoOptimizationFinder<'a, 'tcx> { | |
57 | fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) { | |
58 | let _: Option<_> = try { | |
59 | let target = terminator.kind.as_goto()?; | |
60 | // We only apply this optimization if the last statement is a const assignment | |
61 | let last_statement = self.body.basic_blocks()[location.block].statements.last()?; | |
62 | ||
63 | if let (place, Rvalue::Use(Operand::Constant(_const))) = | |
64 | last_statement.kind.as_assign()? | |
65 | { | |
66 | // We found a constant being assigned to `place`. | |
67 | // Now check that the target of this Goto switches on this place. | |
68 | let target_bb = &self.body.basic_blocks()[target]; | |
69 | ||
70 | // FIXME(simonvandel): We are conservative here when we don't allow | |
71 | // any statements in the target basic block. | |
72 | // This could probably be relaxed to allow `StorageDead`s which could be | |
73 | // copied to the predecessor of this block. | |
74 | if !target_bb.statements.is_empty() { | |
75 | None? | |
76 | } | |
77 | ||
78 | let target_bb_terminator = target_bb.terminator(); | |
79 | let (discr, switch_ty, targets) = target_bb_terminator.kind.as_switch()?; | |
80 | if discr.place() == Some(*place) { | |
81 | // We now know that the Switch matches on the const place, and it is statementless | |
82 | // Now find which value in the Switch matches the const value. | |
83 | let const_value = | |
84 | _const.literal.try_eval_bits(self.tcx, self.param_env, switch_ty)?; | |
85 | let found_value_idx_option = targets | |
86 | .iter() | |
87 | .enumerate() | |
88 | .find(|(_, (value, _))| const_value == *value) | |
89 | .map(|(idx, _)| idx); | |
90 | ||
91 | let target_to_use_in_goto = | |
92 | if let Some(found_value_idx) = found_value_idx_option { | |
93 | targets.iter().nth(found_value_idx).unwrap().1 | |
94 | } else { | |
95 | // If we did not find the const value in values, it must be the otherwise case | |
96 | targets.otherwise() | |
97 | }; | |
98 | ||
99 | self.optimizations.push(OptimizationToApply { | |
100 | bb_with_goto: location.block, | |
101 | target_to_use_in_goto, | |
102 | }); | |
103 | } | |
104 | } | |
105 | Some(()) | |
106 | }; | |
107 | ||
108 | self.super_terminator(terminator, location); | |
109 | } | |
110 | } | |
111 | ||
112 | struct OptimizationToApply { | |
113 | bb_with_goto: BasicBlock, | |
114 | target_to_use_in_goto: BasicBlock, | |
115 | } | |
116 | ||
117 | pub struct ConstGotoOptimizationFinder<'a, 'tcx> { | |
118 | tcx: TyCtxt<'tcx>, | |
119 | body: &'a Body<'tcx>, | |
120 | param_env: ParamEnv<'tcx>, | |
121 | optimizations: Vec<OptimizationToApply>, | |
122 | } |