[MIR] Enhance the SimplifyCfg pass to merge consecutive blocks

src/librustc_mir/transform/simplify_cfg.rs

@@ -21,28 +21,7 @@ impl SimplifyCfg {
         SimplifyCfg
     }
 
-    fn remove_dead_blocks(&self, mir: &mut Mir) {
-        let mut seen = vec![false; mir.basic_blocks.len()];
-
-        // These blocks are always required.
-        seen[START_BLOCK.index()] = true;
-        seen[END_BLOCK.index()] = true;
-        seen[DIVERGE_BLOCK.index()] = true;
-
-        let mut worklist = vec![START_BLOCK];
-        while let Some(bb) = worklist.pop() {
-            for succ in mir.basic_block_data(bb).terminator.successors() {
-                if !seen[succ.index()] {
-                    seen[succ.index()] = true;
-                    worklist.push(*succ);
-                }
-            }
-        }
-
-        util::retain_basic_blocks(mir, &seen);
-    }
-
-    fn remove_goto_chains(&self, mir: &mut Mir) -> bool {
+    fn merge_consecutive_blocks(&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> {
@@ -65,25 +44,70 @@ impl SimplifyCfg {
             Some(target)
         }
 
+        let mut predecessor_map = util::build_predecessor_map(mir);
+
         let mut changed = false;
-        for bb in mir.all_basic_blocks() {
+        let mut seen = vec![false; mir.basic_blocks.len()];
+        let mut worklist = vec![START_BLOCK];
+        while let Some(bb) = worklist.pop() {
             // Temporarily swap out the terminator we're modifying to keep borrowck happy
             let mut terminator = Terminator::Diverge;
             mem::swap(&mut terminator, &mut mir.basic_block_data_mut(bb).terminator);
 
+            // Shortcut chains of empty blocks that just jump from one to the next
             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 *target != new_target {
+                    changed = true;
+                    predecessor_map.remove_predecessor(*target);
+                    predecessor_map.add_predecessor(new_target);
+                    *target = new_target;
+                }
             }
 
-            mir.basic_block_data_mut(bb).terminator = terminator;
+            // See if we can merge the target block into this one
+            match terminator {
+                Terminator::Goto { target } if target.index() > DIVERGE_BLOCK.index() &&
+                        predecessor_map.num_predecessors(target) == 1 => {
+                    changed = true;
+                    let mut other_data = BasicBlockData {
+                        statements: Vec::new(),
+                        terminator: Terminator::Goto { target: target}
+                    };
+                    mem::swap(&mut other_data, mir.basic_block_data_mut(target));
+
+                    // target used to have 1 predecessor (bb), and still has only one (itself)
+                    // All the successors of target have had target replaced by bb in their
+                    // list of predecessors, keeping the number the same.
+
+                    let data = mir.basic_block_data_mut(bb);
+                    data.statements.append(&mut other_data.statements);
+                    mem::swap(&mut data.terminator, &mut other_data.terminator);
+                }
+                _ => mir.basic_block_data_mut(bb).terminator = terminator
+            }
+
+            for succ in mir.basic_block_data(bb).terminator.successors() {
+                if !seen[succ.index()] {
+                    seen[succ.index()] = true;
+                    worklist.push(*succ);
+                }
+            }
         }
 
+        // These blocks must be retained, so mark them seen even if we didn't see them
+        seen[START_BLOCK.index()] = true;
+        seen[END_BLOCK.index()] = true;
+        seen[DIVERGE_BLOCK.index()] = true;
+
+        // Now get rid of all the blocks we never saw
+        util::retain_basic_blocks(mir, &seen);
+
         changed
     }
 
@@ -125,8 +149,7 @@ impl<'tcx> MirPass<'tcx> for SimplifyCfg {
         let mut changed = true;
         while changed {
             changed = self.simplify_branches(mir);
-            changed |= self.remove_goto_chains(mir);
-            self.remove_dead_blocks(mir);
+            changed |= self.merge_consecutive_blocks(mir);
         }
 
         // FIXME: Should probably be moved into some kind of pass manager

src/librustc_mir/transform/util.rs

@@ -9,6 +9,7 @@
 // except according to those terms.
 
 use rustc::mir::repr::*;
+use rustc::mir::visit::*;
 
 /// Update basic block ids in all terminators using the given replacements,
 /// useful e.g. after removal of several basic blocks to update all terminators
@@ -49,3 +50,61 @@ pub fn retain_basic_blocks(mir: &mut Mir, keep: &[bool]) {
 
     update_basic_block_ids(mir, &replacements);
 }
+
+// A simple map to perform quick lookups of the number of predecessors of a
+// BasicBlock. If a block has the same target more than once, for
+// example in a switch, it will appear in the target's predecessor list
+// multiple times. This makes updating the map easier when modifying the graph.
+pub struct PredecessorMap {
+    map: Vec<usize>,
+}
+
+impl PredecessorMap {
+    pub fn new(num_blocks: usize) -> PredecessorMap {
+        PredecessorMap {
+            map: vec![0; num_blocks],
+        }
+    }
+
+    pub fn num_predecessors(&self, block: BasicBlock) -> usize {
+        self.map[block.index()]
+    }
+
+    pub fn add_predecessor(&mut self, block: BasicBlock) {
+        self.map[block.index()] += 1;
+    }
+
+    pub fn remove_predecessor(&mut self, block: BasicBlock) {
+        self.map[block.index()] -= 1;
+    }
+}
+
+struct PredecessorVisitor {
+    predecessor_map: PredecessorMap,
+}
+
+impl PredecessorVisitor {
+    fn new(num_blocks: usize) -> PredecessorVisitor {
+        PredecessorVisitor {
+            predecessor_map: PredecessorMap::new(num_blocks),
+        }
+    }
+}
+
+impl<'tcx> Visitor<'tcx> for PredecessorVisitor {
+    fn visit_mir(&mut self, mir: &Mir<'tcx>) {
+        self.predecessor_map = PredecessorMap::new(mir.basic_blocks.len());
+        self.super_mir(mir);
+    }
+
+    fn visit_branch(&mut self, _source: BasicBlock, target: BasicBlock) {
+        self.predecessor_map.add_predecessor(target);
+    }
+}
+
+pub fn build_predecessor_map(mir: &Mir) -> PredecessorMap {
+    let mut v = PredecessorVisitor::new(mir.basic_blocks.len());
+    v.visit_mir(mir);
+
+    v.predecessor_map
+}