[MIR] Enhance the SimplifyCfg pass to merge consecutive blocks

Author: dotdash

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,71 @@ 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, bb);
+                    predecessor_map.add_predecessor(new_target, bb);
+                    *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.predecessors(target).len() == 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));
+
+                    predecessor_map.replace_predecessor(target, bb, target);
+                    for succ in other_data.terminator.successors() {
+                        predecessor_map.replace_predecessor(*succ, target, bb);
+                    }
+
+                    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 +150,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,70 @@ 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 predecessors of a BasicBlock.
+// Since BasicBlocks usually only have a small number of predecessors, we use a
+// simple vector. Also, 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 allows to update the map more easily when modifying the graph.
+pub struct PredecessorMap {
+    map: Vec<Vec<BasicBlock>>,
+}
+
+impl PredecessorMap {
+    pub fn new(num_blocks: usize) -> PredecessorMap {
+        PredecessorMap {
+            map: vec![Vec::new(); num_blocks],
+        }
+    }
+
+    pub fn predecessors(&self, block: BasicBlock) -> &[BasicBlock] {
+        &self.map[block.index()]
+    }
+
+    pub fn add_predecessor(&mut self, block: BasicBlock, predecessor: BasicBlock) {
+        self.map[block.index()].push(predecessor);
+    }
+
+    pub fn remove_predecessor(&mut self, block: BasicBlock, predecessor: BasicBlock) {
+        let pos = self.map[block.index()].iter().position(|&p| p == predecessor).expect(
+            &format!("{:?} is not registered as a predecessor of {:?}", predecessor, block));
+
+        self.map[block.index()].swap_remove(pos);
+    }
+
+    pub fn replace_predecessor(&mut self, block: BasicBlock, old: BasicBlock, new: BasicBlock) {
+        self.remove_predecessor(block, old);
+        self.add_predecessor(block, new);
+    }
+}
+
+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, source);
+    }
+}
+
+pub fn build_predecessor_map(mir: &Mir) -> PredecessorMap {
+    let mut v = PredecessorVisitor::new(mir.basic_blocks.len());
+    v.visit_mir(mir);
+
+    v.predecessor_map
+}