[MIR] Enhance the SimplifyCfg pass to merge consecutive blocks

dotdash · dotdash · commit 1ebb07ee24ad · 2016-05-11T18:36:49.000+02:00
diff --git a/src/librustc_mir/transform/mod.rs b/src/librustc_mir/transform/mod.rs
@@ -12,6 +12,7 @@ pub mod remove_dead_blocks;
 pub mod simplify_cfg;
 pub mod erase_regions;
 pub mod no_landing_pads;
+pub mod predecessor_map;
 pub mod type_check;
 pub mod break_critical_edges;
 pub mod promote_consts;
diff --git a/src/librustc_mir/transform/predecessor_map.rs b/src/librustc_mir/transform/predecessor_map.rs
@@ -0,0 +1,68 @@
+// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+use rustc::mir::repr::*;
+use rustc::mir::visit::Visitor;
+
+// 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 from_mir(mir: &Mir) -> PredecessorMap {
+        let mut map = PredecessorMap {
+            map: vec![Vec::new(); mir.basic_blocks.len()],
+        };
+
+        PredecessorVisitor { predecessor_map: &mut map }.visit_mir(mir);
+
+        map
+    }
+
+    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);
+    }
+
+    pub fn replace_successor(&mut self, block: BasicBlock, old: BasicBlock, new: BasicBlock) {
+        self.remove_predecessor(old, block);
+        self.add_predecessor(new, block);
+    }
+}
+
+struct PredecessorVisitor<'a> {
+    predecessor_map: &'a mut PredecessorMap,
+}
+
+impl<'a, 'tcx> Visitor<'tcx> for PredecessorVisitor<'a> {
+    fn visit_branch(&mut self, source: BasicBlock, target: BasicBlock) {
+        self.predecessor_map.add_predecessor(target, source);
+    }
+}
diff --git a/src/librustc_mir/transform/simplify_cfg.rs b/src/librustc_mir/transform/simplify_cfg.rs
@@ -8,12 +8,15 @@
 // 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::predecessor_map::PredecessorMap;
 use super::remove_dead_blocks::RemoveDeadBlocks;
 
 pub struct SimplifyCfg;
@@ -22,59 +25,133 @@ 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);
 
-    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;
+        // 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) {
+    let mut predecessor_map = PredecessorMap::from_mir(mir);
+
+    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;
+
+                if let TerminatorKind::Goto { target } = terminator.kind {
+                    // Don't bother trying to merge a block into itself
+                    if target == bb {
+                        break;
+                    }
+
+                    let num_preds = predecessor_map.predecessors(target).len();
+                    let num_insts = mir.basic_block_data(target).statements.len();
+                    match mir.basic_block_data(target).terminator().kind {
+                        _ if num_preds == 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());
+                            }
+
+                            mir.basic_block_data_mut(bb).statements.append(&mut stmts);
+
+                            predecessor_map.replace_predecessor(target, bb, target);
+                            for succ in terminator.successors().iter() {
+                                predecessor_map.replace_predecessor(*succ, target, bb);
                             }
-                            seen.push(next);
-                            target = next;
                         }
-                        _ => break
+                        TerminatorKind::Goto { target: new_target } if num_insts == 0 => {
+                            inner_change = true;
+                            terminator.kind = TerminatorKind::Goto { target: new_target };
+                            predecessor_map.replace_successor(bb, target, new_target);
+                        }
+                        _ => {}
+                    };
+                }
+
+                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;
+                        predecessor_map.replace_successor(bb, *target, new_target);
+                        *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 +183,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 {}
