Implement multiple return terminators optimization

2 files changed, 41 insertions, 0 deletions

diff --git a/compiler/rustc_mir/src/transform/mod.rs b/compiler/rustc_mir/src/transform/mod.rs
index 850cafcf898..67193814a40 100644
--- a/compiler/rustc_mir/src/transform/mod.rs
+++ b/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.
diff --git a/compiler/rustc_mir/src/transform/multiple_return_terminators.rs b/compiler/rustc_mir/src/transform/multiple_return_terminators.rs
new file mode 100644
index 00000000000..3c9c8454f74
--- /dev/null
+++ b/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)
+    }
+}