diff --git a/src/librustc_mir/transform/simplify_cfg.rs b/src/librustc_mir/transform/simplify_cfg.rs index 558276a13a8b8..b4b38506311cb 100644 --- a/src/librustc_mir/transform/simplify_cfg.rs +++ b/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 { @@ -65,25 +44,74 @@ 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); + let mut terminator = mem::replace(&mut mir.basic_block_data_mut(bb).terminator, + Terminator::Diverge); + + // See if we can merge the target block into this one + while let Terminator::Goto { target } = terminator { + if target.index() <= DIVERGE_BLOCK.index() || + predecessor_map.num_predecessors(target) > 1 { + break; + } + + changed = true; + + let mut other_data = mem::replace(mir.basic_block_data_mut(target), BasicBlockData { + statements: Vec::new(), + terminator: Terminator::Goto { target: 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); + terminator = other_data.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; + } } + // Restore the terminator we swapped out for Diverge 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 +153,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 diff --git a/src/librustc_mir/transform/util.rs b/src/librustc_mir/transform/util.rs index 9510269454485..ee06ee0ba41e1 100644 --- a/src/librustc_mir/transform/util.rs +++ b/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, +} + +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 +}