Rollup merge of rust-lang#119699 - cjgillot:simplify-unreachable, r=oli-obk

Merge dead bb pruning and unreachable bb deduplication.

Both routines share the same basic structure: iterate on all bbs to identify work, and then renumber bbs.

We can do both at once.

Author: GuillaumeGomez

compiler/rustc_middle/src/mir/mod.rs

@@ -1316,6 +1316,7 @@ impl<'tcx> BasicBlockData<'tcx> {
     }
 
     /// Does the block have no statements and an unreachable terminator?
+    #[inline]
     pub fn is_empty_unreachable(&self) -> bool {
         self.statements.is_empty() && matches!(self.terminator().kind, TerminatorKind::Unreachable)
     }

compiler/rustc_mir_transform/src/const_goto.rs

@@ -51,7 +51,7 @@ impl<'tcx> MirPass<'tcx> for ConstGoto {
         // if we applied optimizations, we potentially have some cfg to cleanup to
         // make it easier for further passes
         if should_simplify {
-            simplify_cfg(tcx, body);
+            simplify_cfg(body);
             simplify_locals(body, tcx);
         }
     }

compiler/rustc_mir_transform/src/deduplicate_blocks.rs

@@ -25,7 +25,7 @@ impl<'tcx> MirPass<'tcx> for DeduplicateBlocks {
         if has_opts_to_apply {
             let mut opt_applier = OptApplier { tcx, duplicates };
             opt_applier.visit_body(body);
-            simplify_cfg(tcx, body);
+            simplify_cfg(body);
         }
     }
 }

compiler/rustc_mir_transform/src/early_otherwise_branch.rs

@@ -212,7 +212,7 @@ impl<'tcx> MirPass<'tcx> for EarlyOtherwiseBranch {
         // Since this optimization adds new basic blocks and invalidates others,
         // clean up the cfg to make it nicer for other passes
         if should_cleanup {
-            simplify_cfg(tcx, body);
+            simplify_cfg(body);
         }
     }
 }

compiler/rustc_mir_transform/src/inline.rs

@@ -15,7 +15,7 @@ use rustc_target::abi::FieldIdx;
 use rustc_target::spec::abi::Abi;
 
 use crate::cost_checker::CostChecker;
-use crate::simplify::{remove_dead_blocks, CfgSimplifier};
+use crate::simplify::simplify_cfg;
 use crate::util;
 use std::iter;
 use std::ops::{Range, RangeFrom};
@@ -56,8 +56,7 @@ impl<'tcx> MirPass<'tcx> for Inline {
         let _guard = span.enter();
         if inline(tcx, body) {
             debug!("running simplify cfg on {:?}", body.source);
-            CfgSimplifier::new(body).simplify();
-            remove_dead_blocks(body);
+            simplify_cfg(body);
             deref_finder(tcx, body);
         }
     }

compiler/rustc_mir_transform/src/match_branches.rs

@@ -174,7 +174,7 @@ impl<'tcx> MirPass<'tcx> for MatchBranchSimplification {
         }
 
         if should_cleanup {
-            simplify_cfg(tcx, body);
+            simplify_cfg(body);
         }
     }
 }

compiler/rustc_mir_transform/src/remove_unneeded_drops.rs

@@ -38,7 +38,7 @@ impl<'tcx> MirPass<'tcx> for RemoveUnneededDrops {
         // if we applied optimizations, we potentially have some cfg to cleanup to
         // make it easier for further passes
         if should_simplify {
-            simplify_cfg(tcx, body);
+            simplify_cfg(body);
         }
     }
 }

compiler/rustc_mir_transform/src/separate_const_switch.rs

@@ -50,11 +50,11 @@ impl<'tcx> MirPass<'tcx> for SeparateConstSwitch {
         sess.mir_opt_level() >= 2 && sess.opts.unstable_opts.unsound_mir_opts
     }
 
-    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+    fn run_pass(&self, _: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         // If execution did something, applying a simplification layer
         // helps later passes optimize the copy away.
         if separate_const_switch(body) > 0 {
-            super::simplify::simplify_cfg(tcx, body);
+            super::simplify::simplify_cfg(body);
         }
     }
 }

compiler/rustc_mir_transform/src/simplify.rs

@@ -27,7 +27,6 @@
 //! naively generate still contains the `_a = ()` write in the unreachable block "after" the
 //! return.
 
-use rustc_data_structures::fx::FxIndexSet;
 use rustc_index::{Idx, IndexSlice, IndexVec};
 use rustc_middle::mir::visit::{MutVisitor, MutatingUseContext, PlaceContext, Visitor};
 use rustc_middle::mir::*;
@@ -62,9 +61,8 @@ impl SimplifyCfg {
     }
 }
 
-pub fn simplify_cfg<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+pub(crate) fn simplify_cfg(body: &mut Body<'_>) {
     CfgSimplifier::new(body).simplify();
-    remove_duplicate_unreachable_blocks(tcx, body);
     remove_dead_blocks(body);
 
     // FIXME: Should probably be moved into some kind of pass manager
@@ -76,9 +74,9 @@ impl<'tcx> MirPass<'tcx> for SimplifyCfg {
         self.name()
     }
 
-    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+    fn run_pass(&self, _: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         debug!("SimplifyCfg({:?}) - simplifying {:?}", self.name(), body.source);
-        simplify_cfg(tcx, body);
+        simplify_cfg(body);
     }
 }
 
@@ -289,55 +287,25 @@ pub fn simplify_duplicate_switch_targets(terminator: &mut Terminator<'_>) {
     }
 }
 
-pub fn remove_duplicate_unreachable_blocks<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-    struct OptApplier<'tcx> {
-        tcx: TyCtxt<'tcx>,
-        duplicates: FxIndexSet<BasicBlock>,
-    }
-
-    impl<'tcx> MutVisitor<'tcx> for OptApplier<'tcx> {
-        fn tcx(&self) -> TyCtxt<'tcx> {
-            self.tcx
-        }
-
-        fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) {
-            for target in terminator.successors_mut() {
-                // We don't have to check whether `target` is a cleanup block, because have
-                // entirely excluded cleanup blocks in building the set of duplicates.
-                if self.duplicates.contains(target) {
-                    *target = self.duplicates[0];
-                }
-            }
-
-            simplify_duplicate_switch_targets(terminator);
-
-            self.super_terminator(terminator, location);
-        }
-    }
+pub(crate) fn remove_dead_blocks(body: &mut Body<'_>) {
+    let should_deduplicate_unreachable = |bbdata: &BasicBlockData<'_>| {
+        // CfgSimplifier::simplify leaves behind some unreachable basic blocks without a
+        // terminator. Those blocks will be deleted by remove_dead_blocks, but we run just
+        // before then so we need to handle missing terminators.
+        // We also need to prevent confusing cleanup and non-cleanup blocks. In practice we
+        // don't emit empty unreachable cleanup blocks, so this simple check suffices.
+        bbdata.terminator.is_some() && bbdata.is_empty_unreachable() && !bbdata.is_cleanup
+    };
 
-    let unreachable_blocks = body
+    let reachable = traversal::reachable_as_bitset(body);
+    let empty_unreachable_blocks = body
         .basic_blocks
         .iter_enumerated()
-        .filter(|(_, bb)| {
-            // CfgSimplifier::simplify leaves behind some unreachable basic blocks without a
-            // terminator. Those blocks will be deleted by remove_dead_blocks, but we run just
-            // before then so we need to handle missing terminators.
-            // We also need to prevent confusing cleanup and non-cleanup blocks. In practice we
-            // don't emit empty unreachable cleanup blocks, so this simple check suffices.
-            bb.terminator.is_some() && bb.is_empty_unreachable() && !bb.is_cleanup
-        })
-        .map(|(block, _)| block)
-        .collect::<FxIndexSet<_>>();
-
-    if unreachable_blocks.len() > 1 {
-        OptApplier { tcx, duplicates: unreachable_blocks }.visit_body(body);
-    }
-}
+        .filter(|(bb, bbdata)| should_deduplicate_unreachable(bbdata) && reachable.contains(*bb))
+        .count();
 
-pub fn remove_dead_blocks(body: &mut Body<'_>) {
-    let reachable = traversal::reachable_as_bitset(body);
     let num_blocks = body.basic_blocks.len();
-    if num_blocks == reachable.count() {
+    if num_blocks == reachable.count() && empty_unreachable_blocks <= 1 {
         return;
     }
 
@@ -346,14 +314,28 @@ pub fn remove_dead_blocks(body: &mut Body<'_>) {
     let mut replacements: Vec<_> = (0..num_blocks).map(BasicBlock::new).collect();
     let mut orig_index = 0;
     let mut used_index = 0;
-    basic_blocks.raw.retain(|_| {
-        let keep = reachable.contains(BasicBlock::new(orig_index));
-        if keep {
-            replacements[orig_index] = BasicBlock::new(used_index);
-            used_index += 1;
+    let mut kept_unreachable = None;
+    basic_blocks.raw.retain(|bbdata| {
+        let orig_bb = BasicBlock::new(orig_index);
+        if !reachable.contains(orig_bb) {
+            orig_index += 1;
+            return false;
+        }
+
+        let used_bb = BasicBlock::new(used_index);
+        if should_deduplicate_unreachable(bbdata) {
+            let kept_unreachable = *kept_unreachable.get_or_insert(used_bb);
+            if kept_unreachable != used_bb {
+                replacements[orig_index] = kept_unreachable;
+                orig_index += 1;
+                return false;
+            }
         }
+
+        replacements[orig_index] = used_bb;
+        used_index += 1;
         orig_index += 1;
-        keep
+        true
     });
 
     for block in basic_blocks {

tests/mir-opt/building/async_await.b-{closure#0}.coroutine_resume.0.mir

@@ -108,7 +108,7 @@ fn b::{closure#0}(_1: Pin<&mut {async fn body@$DIR/async_await.rs:15:18: 18:2}>,
 
     bb0: {
         _39 = discriminant((*(_1.0: &mut {async fn body@$DIR/async_await.rs:15:18: 18:2})));
-        switchInt(move _39) -> [0: bb1, 1: bb29, 3: bb27, 4: bb28, otherwise: bb30];
+        switchInt(move _39) -> [0: bb1, 1: bb29, 3: bb27, 4: bb28, otherwise: bb9];
     }
 
     bb1: {
@@ -345,8 +345,4 @@ fn b::{closure#0}(_1: Pin<&mut {async fn body@$DIR/async_await.rs:15:18: 18:2}>,
     bb29: {
         assert(const false, "`async fn` resumed after completion") -> [success: bb29, unwind unreachable];
     }
-
-    bb30: {
-        unreachable;
-    }
 }