diff --git a/src/librustc_mir/transform/simplify_cfg.rs b/src/librustc_mir/transform/simplify_cfg.rs index 71dd2f077fec2..e9ce04001f742 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,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 MirPass 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 e45cfa8395472..ff21ac0479ac4 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 repr::*; +use 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>, +} + +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 +}