forked from rust-lang/rust
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
New mir-opt pass to simplify gotos with const values
Fixes rust-lang#77355
- Loading branch information
1 parent
0dce3f6
commit f24c0ea
Showing
8 changed files
with
260 additions
and
117 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,129 @@ | ||
| //! This pass optimizes the following sequence | ||
| //! ```rust | ||
| //! bb2: { | ||
| //! _2 = const true; | ||
| //! goto -> bb3; | ||
| //! } | ||
| //! | ||
| //! bb3: { | ||
| //! switchInt(_2) -> [false: bb4, otherwise: bb5]; | ||
| //! } | ||
| //! ``` | ||
| //! into | ||
| //! ```rust | ||
| //! bb2: { | ||
| //! _2 = const true; | ||
| //! goto -> bb5; | ||
| //! } | ||
| //! ``` | ||
| use crate::transform::MirPass; | ||
| use rustc_middle::mir::*; | ||
| use rustc_middle::ty::TyCtxt; | ||
| use rustc_middle::{mir::visit::Visitor, ty::ParamEnv}; | ||
|
|
||
| use super::simplify::{simplify_cfg, simplify_locals}; | ||
|
|
||
| pub struct ConstGoto; | ||
|
|
||
| impl<'tcx> MirPass<'tcx> for ConstGoto { | ||
| fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { | ||
| trace!("Running ConstGoto on {:?}", body.source); | ||
| let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id()); | ||
| let mut opt_finder = | ||
| ConstGotoOptimizationFinder { tcx, body, optimizations: vec![], param_env }; | ||
| opt_finder.visit_body(body); | ||
| let should_simplify = !opt_finder.optimizations.is_empty(); | ||
| for opt in opt_finder.optimizations { | ||
| let terminator = body.basic_blocks_mut()[opt.bb_with_goto].terminator_mut(); | ||
| let new_goto = TerminatorKind::Goto { target: opt.target_to_use_in_goto }; | ||
| debug!("SUCCESS: replacing `{:?}` with `{:?}`", terminator.kind, new_goto); | ||
| terminator.kind = new_goto; | ||
| } | ||
|
|
||
| // if we applied optimizations, we potentially have some cfg to cleanup to | ||
| // make it easier for further passes | ||
| if should_simplify { | ||
| simplify_cfg(body); | ||
| simplify_locals(body, tcx); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| impl<'a, 'tcx> Visitor<'tcx> for ConstGotoOptimizationFinder<'a, 'tcx> { | ||
| fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) { | ||
| let mut bailer = || { | ||
| match terminator.kind { | ||
| TerminatorKind::Goto { target } => { | ||
| // We only apply this optimization if the last statement is a const assignment | ||
| let last_statement = | ||
| self.body.basic_blocks()[location.block].statements.last()?; | ||
|
|
||
| match &last_statement.kind { | ||
| StatementKind::Assign(box (place, Rvalue::Use(op))) => { | ||
| let _const = op.constant()?; | ||
| // We found a constant being assigned to `place`. | ||
| // Now check that the target of this Goto switches on this place. | ||
| let target_bb = &self.body.basic_blocks()[target]; | ||
| if !target_bb.statements.is_empty() { | ||
| return None; | ||
| } | ||
|
|
||
| let target_bb_terminator = target_bb.terminator(); | ||
| match &target_bb_terminator.kind { | ||
| TerminatorKind::SwitchInt { discr, switch_ty, targets } | ||
| if discr.place() == Some(*place) => | ||
| { | ||
| // We now know that the Switch matches on the const place, and it is statementless | ||
| // Now find which value in the Switch matches the const value. | ||
| let const_value = _const.literal.eval_bits( | ||
| self.tcx, | ||
| self.param_env, | ||
| switch_ty, | ||
| ); | ||
| let found_value_idx_option = targets | ||
| .iter() | ||
| .enumerate() | ||
| .find(|(_, x)| const_value == x.0) | ||
| .map(|(idx, _)| idx); | ||
|
|
||
| let target_to_use_in_goto = | ||
| if let Some(found_value_idx) = found_value_idx_option { | ||
| targets.iter().nth(found_value_idx).unwrap().1 | ||
| } else { | ||
| // If we did not find the const value in values, it must be the otherwise case | ||
| targets.otherwise() | ||
| }; | ||
|
|
||
| self.optimizations.push(OptimizationToApply { | ||
| bb_with_goto: location.block, | ||
| target_to_use_in_goto, | ||
| }); | ||
| } | ||
| _ => {} | ||
| } | ||
| } | ||
| _ => {} | ||
| } | ||
| } | ||
| _ => {} | ||
| } | ||
| return Some(()); | ||
| }; | ||
| let _ = bailer(); | ||
|
|
||
| self.super_terminator(terminator, location); | ||
| } | ||
| } | ||
|
|
||
| struct OptimizationToApply { | ||
| bb_with_goto: BasicBlock, | ||
| target_to_use_in_goto: BasicBlock, | ||
| } | ||
|
|
||
| pub struct ConstGotoOptimizationFinder<'a, 'tcx> { | ||
| tcx: TyCtxt<'tcx>, | ||
| body: &'a Body<'tcx>, | ||
| param_env: ParamEnv<'tcx>, | ||
| optimizations: Vec<OptimizationToApply>, | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
52 changes: 52 additions & 0 deletions
52
src/test/mir-opt/const_goto.issue_77355_opt.ConstGoto.diff
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,52 @@ | ||
| - // MIR for `issue_77355_opt` before ConstGoto | ||
| + // MIR for `issue_77355_opt` after ConstGoto | ||
|
|
||
| fn issue_77355_opt(_1: Foo) -> u64 { | ||
| debug num => _1; // in scope 0 at $DIR/const_goto.rs:11:20: 11:23 | ||
| let mut _0: u64; // return place in scope 0 at $DIR/const_goto.rs:11:33: 11:36 | ||
| - let mut _2: bool; // in scope 0 at $SRC_DIR/core/src/macros/mod.rs:LL:COL | ||
| - let mut _3: isize; // in scope 0 at $DIR/const_goto.rs:12:22: 12:28 | ||
| + let mut _2: isize; // in scope 0 at $DIR/const_goto.rs:12:22: 12:28 | ||
|
|
||
| bb0: { | ||
| - StorageLive(_2); // scope 0 at $SRC_DIR/core/src/macros/mod.rs:LL:COL | ||
| - _3 = discriminant(_1); // scope 0 at $DIR/const_goto.rs:12:22: 12:28 | ||
| - switchInt(move _3) -> [1_isize: bb2, 2_isize: bb2, otherwise: bb1]; // scope 0 at $DIR/const_goto.rs:12:22: 12:28 | ||
| + _2 = discriminant(_1); // scope 0 at $DIR/const_goto.rs:12:22: 12:28 | ||
| + switchInt(move _2) -> [1_isize: bb2, 2_isize: bb2, otherwise: bb1]; // scope 0 at $DIR/const_goto.rs:12:22: 12:28 | ||
| } | ||
|
|
||
| bb1: { | ||
| - _2 = const false; // scope 0 at $SRC_DIR/core/src/macros/mod.rs:LL:COL | ||
| - goto -> bb3; // scope 0 at $SRC_DIR/core/src/macros/mod.rs:LL:COL | ||
| - } | ||
| - | ||
| - bb2: { | ||
| - _2 = const true; // scope 0 at $SRC_DIR/core/src/macros/mod.rs:LL:COL | ||
| - goto -> bb3; // scope 0 at $SRC_DIR/core/src/macros/mod.rs:LL:COL | ||
| - } | ||
| - | ||
| - bb3: { | ||
| - switchInt(_2) -> [false: bb4, otherwise: bb5]; // scope 0 at $DIR/const_goto.rs:12:5: 12:57 | ||
| - } | ||
| - | ||
| - bb4: { | ||
| _0 = const 42_u64; // scope 0 at $DIR/const_goto.rs:12:53: 12:55 | ||
| - goto -> bb6; // scope 0 at $DIR/const_goto.rs:12:5: 12:57 | ||
| + goto -> bb3; // scope 0 at $DIR/const_goto.rs:12:5: 12:57 | ||
| } | ||
|
|
||
| - bb5: { | ||
| + bb2: { | ||
| _0 = const 23_u64; // scope 0 at $DIR/const_goto.rs:12:41: 12:43 | ||
| - goto -> bb6; // scope 0 at $DIR/const_goto.rs:12:5: 12:57 | ||
| + goto -> bb3; // scope 0 at $DIR/const_goto.rs:12:5: 12:57 | ||
| } | ||
|
|
||
| - bb6: { | ||
| - StorageDead(_2); // scope 0 at $DIR/const_goto.rs:13:1: 13:2 | ||
| + bb3: { | ||
| return; // scope 0 at $DIR/const_goto.rs:13:2: 13:2 | ||
| } | ||
| } | ||
|
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,16 @@ | ||
| pub enum Foo { | ||
| A, | ||
| B, | ||
| C, | ||
| D, | ||
| E, | ||
| F, | ||
| } | ||
|
|
||
| // EMIT_MIR const_goto.issue_77355_opt.ConstGoto.diff | ||
| fn issue_77355_opt(num: Foo) -> u64 { | ||
| if matches!(num, Foo::B | Foo::C) { 23 } else { 42 } | ||
| } | ||
| fn main() { | ||
| issue_77355_opt(Foo::A); | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.