Rollup merge of #33552 - dotdash:scfg, r=luqmana

[MIR] Enhance the SimplifyCfg pass to merge consecutive blocks

Updated from #30238, including the changes suggested by @Aatch.

1 files changed, 124 insertions, 59 deletions

diff --git a/src/librustc_mir/transform/simplify_cfg.rs b/src/librustc_mir/transform/simplify_cfg.rs
index fa897384a54..526157a49c7 100644
--- a/src/librustc_mir/transform/simplify_cfg.rs
+++ b/src/librustc_mir/transform/simplify_cfg.rs
@@ -8,73 +8,155 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
+use rustc_data_structures::bitvec::BitVector;
 use rustc::middle::const_val::ConstVal;
 use rustc::ty::TyCtxt;
 use rustc::mir::repr::*;
 use rustc::mir::transform::{MirPass, MirSource, Pass};
 use pretty;
+use std::mem;
 
 use super::remove_dead_blocks::RemoveDeadBlocks;
 
+use traversal;
+
 pub struct SimplifyCfg;
 
 impl SimplifyCfg {
     pub fn new() -> SimplifyCfg {
         SimplifyCfg
     }
+}
+
+impl<'tcx> MirPass<'tcx> for SimplifyCfg {
+    fn run_pass<'a>(&mut self, tcx: TyCtxt<'a, 'tcx, 'tcx>, src: MirSource, mir: &mut Mir<'tcx>) {
+        simplify_branches(mir);
+        RemoveDeadBlocks.run_pass(tcx, src, mir);
+        merge_consecutive_blocks(mir);
+        RemoveDeadBlocks.run_pass(tcx, src, mir);
+        pretty::dump_mir(tcx, "simplify_cfg", &0, src, mir, None);
+
+        // FIXME: Should probably be moved into some kind of pass manager
+        mir.basic_blocks.shrink_to_fit();
+    }
+}
+
+impl Pass for SimplifyCfg {}
+
+fn merge_consecutive_blocks(mir: &mut Mir) {
+    // Build the precedecessor map for the MIR
+    let mut pred_count = vec![0u32; mir.basic_blocks.len()];
+    for (_, data) in traversal::preorder(mir) {
+        if let Some(ref term) = data.terminator {
+            for &tgt in term.successors().iter() {
+                pred_count[tgt.index()] += 1;
+            }
+        }
+    }
+
+    loop {
+        let mut changed = false;
+        let mut seen = BitVector::new(mir.basic_blocks.len());
+        let mut worklist = vec![START_BLOCK];
+        while let Some(bb) = worklist.pop() {
+            // Temporarily take ownership of the terminator we're modifying to keep borrowck happy
+            let mut terminator = mir.basic_block_data_mut(bb).terminator.take()
+                .expect("invalid terminator state");
+
+            // See if we can merge the target block into this one
+            loop {
+                let mut inner_change = false;
 
-    fn remove_goto_chains(&self, mir: &mut Mir) -> bool {
-        // Find the target at the end of the jump chain, return None if there is a loop
-        fn final_target(mir: &Mir, mut target: BasicBlock) -> Option<BasicBlock> {
-            // Keep track of already seen blocks to detect loops
-            let mut seen: Vec<BasicBlock> = Vec::with_capacity(8);
-
-            while mir.basic_block_data(target).statements.is_empty() {
-                // NB -- terminator may have been swapped with `None`
-                // below, in which case we have a cycle and just want
-                // to stop
-                if let Some(ref terminator) = mir.basic_block_data(target).terminator {
-                    match terminator.kind {
-                        TerminatorKind::Goto { target: next } => {
-                            if seen.contains(&next) {
-                                return None;
+                if let TerminatorKind::Goto { target } = terminator.kind {
+                    // Don't bother trying to merge a block into itself
+                    if target == bb {
+                        break;
+                    }
+
+                    let num_insts = mir.basic_block_data(target).statements.len();
+                    match mir.basic_block_data(target).terminator().kind {
+                        TerminatorKind::Goto { target: new_target } if num_insts == 0 => {
+                            inner_change = true;
+                            terminator.kind = TerminatorKind::Goto { target: new_target };
+                            pred_count[target.index()] -= 1;
+                            pred_count[new_target.index()] += 1;
+                        }
+                        _ if pred_count[target.index()] == 1 => {
+                            inner_change = true;
+                            let mut stmts = Vec::new();
+                            {
+                                let target_data = mir.basic_block_data_mut(target);
+                                mem::swap(&mut stmts, &mut target_data.statements);
+                                mem::swap(&mut terminator, target_data.terminator_mut());
                             }
-                            seen.push(next);
-                            target = next;
+
+                            mir.basic_block_data_mut(bb).statements.append(&mut stmts);
                         }
-                        _ => break
+                        _ => {}
+                    };
+                }
+
+                for target in terminator.successors_mut() {
+                    let new_target = match final_target(mir, *target) {
+                        Some(new_target) => new_target,
+                        None if mir.basic_block_data(bb).statements.is_empty() => bb,
+                        None => continue
+                    };
+                    if *target != new_target {
+                        inner_change = true;
+                        pred_count[target.index()] -= 1;
+                        pred_count[new_target.index()] += 1;
+                        *target = new_target;
                     }
-                } else {
-                    break
+                }
+
+                changed |= inner_change;
+                if !inner_change {
+                    break;
                 }
             }
 
-            Some(target)
+            mir.basic_block_data_mut(bb).terminator = Some(terminator);
+
+            for succ in mir.basic_block_data(bb).terminator().successors().iter() {
+                if seen.insert(succ.index()) {
+                    worklist.push(*succ);
+                }
+            }
         }
 
-        let mut changed = false;
-        for bb in mir.all_basic_blocks() {
-            // Temporarily take ownership of the terminator we're modifying to keep borrowck happy
-            let mut terminator = mir.basic_block_data_mut(bb).terminator.take()
-                                    .expect("invalid terminator state");
-
-            debug!("remove_goto_chains: bb={:?} terminator={:?}", bb, terminator);
-
-            for target in terminator.successors_mut() {
-                let new_target = match final_target(mir, *target) {
-                    Some(new_target) => new_target,
-                    None if mir.basic_block_data(bb).statements.is_empty() => bb,
-                    None => continue
-                };
-                changed |= *target != new_target;
-                *target = new_target;
+        if !changed {
+            break;
+        }
+    }
+}
+
+// Find the target at the end of the jump chain, return None if there is a loop
+fn final_target(mir: &Mir, mut target: BasicBlock) -> Option<BasicBlock> {
+    // Keep track of already seen blocks to detect loops
+    let mut seen: Vec<BasicBlock> = Vec::with_capacity(8);
+
+    while mir.basic_block_data(target).statements.is_empty() {
+        // NB -- terminator may have been swapped with `None` in
+        // merge_consecutive_blocks, in which case we have a cycle and just want
+        // to stop
+        match mir.basic_block_data(target).terminator {
+            Some(Terminator { kind: TerminatorKind::Goto { target: next }, .. }) =>  {
+                if seen.contains(&next) {
+                    return None;
+                }
+                seen.push(next);
+                target = next;
             }
-            mir.basic_block_data_mut(bb).terminator = Some(terminator);
+            _ => break
         }
-        changed
     }
 
-    fn simplify_branches(&self, mir: &mut Mir) -> bool {
+    Some(target)
+}
+
+fn simplify_branches(mir: &mut Mir) {
+    loop {
         let mut changed = false;
 
         for bb in mir.all_basic_blocks() {
@@ -106,25 +188,8 @@ impl SimplifyCfg {
             }
         }
 
-        changed
-    }
-}
-
-impl<'tcx> MirPass<'tcx> for SimplifyCfg {
-    fn run_pass<'a>(&mut self, tcx: TyCtxt<'a, 'tcx, 'tcx>,
-                    src: MirSource, mir: &mut Mir<'tcx>) {
-        let mut counter = 0;
-        let mut changed = true;
-        while changed {
-            pretty::dump_mir(tcx, "simplify_cfg", &counter, src, mir, None);
-            counter += 1;
-            changed = self.simplify_branches(mir);
-            changed |= self.remove_goto_chains(mir);
-            RemoveDeadBlocks.run_pass(tcx, src, mir);
+        if !changed {
+            break;
         }
-        // FIXME: Should probably be moved into some kind of pass manager
-        mir.basic_blocks.shrink_to_fit();
     }
 }
-
-impl Pass for SimplifyCfg {}