Implement multiple return terminators optimization

Author: alarsyo

compiler/rustc_mir/src/transform/mod.rs

@@ -33,6 +33,7 @@ pub mod inline;
 pub mod instcombine;
 pub mod instrument_coverage;
 pub mod match_branches;
+pub mod multiple_return_terminators;
 pub mod no_landing_pads;
 pub mod nrvo;
 pub mod promote_consts;
@@ -464,6 +465,7 @@ fn run_optimization_passes<'tcx>(
  &remove_unneeded_drops::RemoveUnneededDrops,
  &match_branches::MatchBranchSimplification,
  // inst combine is after MatchBranchSimplification to clean up Ne(_1, false)
+ &multiple_return_terminators::MultipleReturnTerminators,
  &instcombine::InstCombine,
  &const_prop::ConstProp,
  &simplify_branches::SimplifyBranches::new("after-const-prop"),
@@ -478,6 +480,7 @@ fn run_optimization_passes<'tcx>(
  &simplify::SimplifyCfg::new("final"),
  &nrvo::RenameReturnPlace,
  &simplify::SimplifyLocals,
+ &multiple_return_terminators::MultipleReturnTerminators,
  ];
 
  // Optimizations to run even if mir optimizations have been disabled.

compiler/rustc_mir/src/transform/multiple_return_terminators.rs

@@ -0,0 +1,38 @@
+//! This pass removes jumps to basic blocks containing only a return, and replaces them with a
+//! return instead.
+
+use crate::transform::{simplify, MirPass, MirSource};
+use rustc_index::bit_set::BitSet;
+use rustc_middle::mir::*;
+use rustc_middle::ty::TyCtxt;
+
+pub struct MultipleReturnTerminators;
+
+impl<'tcx> MirPass<'tcx> for MultipleReturnTerminators {
+ fn run_pass(&self, tcx: TyCtxt<'tcx>, _: MirSource<'tcx>, body: &mut Body<'tcx>) {
+ if tcx.sess.opts.debugging_opts.mir_opt_level < 3 {
+ return;
+ }
+
+ // find basic blocks with no statement and a return terminator
+ let mut bbs_simple_returns = BitSet::new_empty(body.basic_blocks().len());
+ let bbs = body.basic_blocks_mut();
+ for idx in bbs.indices() {
+ if bbs[idx].statements.is_empty()
+ && bbs[idx].terminator().kind == TerminatorKind::Return
+ {
+ bbs_simple_returns.insert(idx);
+ }
+ }
+
+ for bb in bbs {
+ if let TerminatorKind::Goto { target } = bb.terminator().kind {
+ if bbs_simple_returns.contains(target) {
+ bb.terminator_mut().kind = TerminatorKind::Return;
+ }
+ }
+ }
+
+ simplify::remove_dead_blocks(body)
+ }
+}