Add and use a break critical edges transform

This is a fairly standard transform that inserts blocks along critical
edges so code can be inserted along the edge without it affecting other
edges. The main difference is that it considers a Drop or Call
terminator that would require an `invoke` instruction in LLVM a critical
edge. This is because we can't actually insert code after an invoke, so
it ends up looking similar to a critical edge anyway.

The transform is run just before translation right now.

5 files changed, 227 insertions, 42 deletions

diff --git a/src/librustc_driver/driver.rs b/src/librustc_driver/driver.rs
index 468dc7b12c1..03f6ce355b8 100644
--- a/src/librustc_driver/driver.rs
+++ b/src/librustc_driver/driver.rs
@@ -932,7 +932,7 @@ pub fn phase_3_run_analysis_passes<'tcx, F, R>(sess: &'tcx Session,
 
 /// Run the translation phase to LLVM, after which the AST and analysis can
 pub fn phase_4_translate_to_llvm<'tcx>(tcx: &TyCtxt<'tcx>,
-                                       mir_map: MirMap<'tcx>,
+                                       mut mir_map: MirMap<'tcx>,
                                        analysis: ty::CrateAnalysis)
                                        -> trans::CrateTranslation {
     let time_passes = tcx.sess.time_passes();
@@ -941,6 +941,13 @@ pub fn phase_4_translate_to_llvm<'tcx>(tcx: &TyCtxt<'tcx>,
          "resolving dependency formats",
          || dependency_format::calculate(&tcx.sess));
 
+    time(time_passes,
+         "erasing regions from MIR",
+         || mir::transform::erase_regions::erase_regions(tcx, &mut mir_map));
+
+    time(time_passes, "breaking critical edges",
+         || mir::transform::break_critical_edges::break_critical_edges(&mut mir_map));
+
     // Option dance to work around the lack of stack once closures.
     time(time_passes,
          "translation",
diff --git a/src/librustc_mir/transform/break_critical_edges.rs b/src/librustc_mir/transform/break_critical_edges.rs
new file mode 100644
index 00000000000..5e64da0ec90
--- /dev/null
+++ b/src/librustc_mir/transform/break_critical_edges.rs
@@ -0,0 +1,180 @@
+// 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 std::mem;
+
+use rustc_back::slice;
+use rustc::mir::repr::*;
+use rustc::mir::mir_map::MirMap;
+
+use traversal;
+
+/**
+ * Breaks critical edges in the MIR.
+ *
+ * Critical edges are edges that are neither the only edge leaving a
+ * block, nor the only edge entering one.
+ *
+ * When you want something to happen "along" an edge, you can either
+ * do at the end of the predecessor block, or at the start of the
+ * successor block. Critical edges have to be broken in order to prevent
+ * "edge actions" from affecting other edges.
+ *
+ * This function will break those edges by inserting new blocks along them.
+ *
+ * A special case is Drop and Call terminators with unwind/cleanup successors,
+ * They use `invoke` in LLVM, which terminates a block, meaning that code cannot
+ * be inserted after them, so even if an edge is the only edge leaving a block
+ * like that, we still insert blocks if the edge is one of many entering the
+ * target.
+ *
+ * NOTE: Simplify CFG will happily undo most of the work this pass does.
+ *
+ */
+pub fn break_critical_edges<'tcx>(mir_map: &mut MirMap<'tcx>) {
+    for (_, mir) in &mut mir_map.map {
+        break_critical_edges_fn(mir);
+    }
+}
+
+/*
+ * Predecessor map for tracking the predecessors of a block
+ */
+struct PredMap {
+    preds: Vec<BlockPredecessors>
+}
+
+/**
+ * Most blocks only have one predecessor, so we can cut down on
+ * some allocation by not using Vec until we have more than one.
+ */
+#[derive(Clone)]
+enum BlockPredecessors {
+    None,
+    One(BasicBlock),
+    Some(Vec<BasicBlock>)
+}
+
+impl PredMap {
+    pub fn new(n: usize) -> PredMap {
+        let preds = vec![BlockPredecessors::None; n];
+
+        PredMap {
+            preds: preds
+        }
+    }
+
+    fn ensure_len(&mut self, bb: BasicBlock) {
+        let idx = bb.index();
+        while self.preds.len() <= idx {
+            self.preds.push(BlockPredecessors::None);
+        }
+    }
+
+    pub fn add_pred(&mut self, target: BasicBlock, pred: BasicBlock) {
+        self.ensure_len(target);
+
+        let preds = mem::replace(&mut self.preds[target.index()], BlockPredecessors::None);
+        match preds {
+            BlockPredecessors::None => {
+                self.preds[target.index()] = BlockPredecessors::One(pred);
+            }
+            BlockPredecessors::One(bb) => {
+                self.preds[target.index()] = BlockPredecessors::Some(vec![bb, pred]);
+            }
+            BlockPredecessors::Some(mut preds) => {
+                preds.push(pred);
+                self.preds[target.index()] = BlockPredecessors::Some(preds);
+            }
+        }
+    }
+
+    pub fn remove_pred(&mut self, target: BasicBlock, pred: BasicBlock) {
+        self.ensure_len(target);
+
+        let preds = mem::replace(&mut self.preds[target.index()], BlockPredecessors::None);
+        match preds {
+            BlockPredecessors::None => {}
+            BlockPredecessors::One(bb) if bb == pred => {}
+
+            BlockPredecessors::One(bb) => {
+                self.preds[target.index()] = BlockPredecessors::One(bb);
+            }
+
+            BlockPredecessors::Some(mut preds) => {
+                preds.retain(|&bb| bb != pred);
+                self.preds[target.index()] = BlockPredecessors::Some(preds);
+            }
+        }
+    }
+
+    pub fn get_preds(&self, bb: BasicBlock) -> &[BasicBlock] {
+        match self.preds[bb.index()] {
+            BlockPredecessors::None => &[],
+            BlockPredecessors::One(ref bb) => slice::ref_slice(bb),
+            BlockPredecessors::Some(ref bbs) => &bbs[..]
+        }
+    }
+}
+
+
+fn break_critical_edges_fn(mir: &mut Mir) {
+    let mut pred_map = PredMap::new(mir.basic_blocks.len());
+
+    // Build the precedecessor map for the MIR
+    for (pred, data) in traversal::preorder(mir) {
+        if let Some(ref term) = data.terminator {
+            for &tgt in term.successors().iter() {
+                pred_map.add_pred(tgt, pred);
+            }
+        }
+    }
+
+    // We need a place to store the new blocks generated
+    let mut new_blocks = Vec::new();
+
+    let bbs = mir.all_basic_blocks();
+    let cur_len = mir.basic_blocks.len();
+
+    for &bb in &bbs {
+        let data = mir.basic_block_data_mut(bb);
+
+        if let Some(ref mut term) = data.terminator {
+            let is_invoke = term_is_invoke(term);
+            let succs = term.successors_mut();
+            if succs.len() > 1 || (succs.len() > 0 && is_invoke) {
+                for tgt in succs {
+                    let num_preds = pred_map.get_preds(*tgt).len();
+                    if num_preds > 1 {
+                        // It's a critical edge, break it
+                        let goto = Terminator::Goto { target: *tgt };
+                        let data = BasicBlockData::new(Some(goto));
+                        // Get the index it will be when inserted into the MIR
+                        let idx = cur_len + new_blocks.len();
+                        new_blocks.push(data);
+                        *tgt = BasicBlock::new(idx);
+                    }
+                }
+            }
+        }
+    }
+
+    debug!("Broke {} N edges", new_blocks.len());
+
+    mir.basic_blocks.extend_from_slice(&new_blocks);
+}
+
+// Returns true if the terminator would use an invoke in LLVM.
+fn term_is_invoke(term: &Terminator) -> bool {
+    match *term {
+        Terminator::Call { cleanup: Some(_), .. } |
+        Terminator::Drop { unwind: Some(_), .. } => true,
+        _ => false
+    }
+}
diff --git a/src/librustc_mir/transform/mod.rs b/src/librustc_mir/transform/mod.rs
index 57690caeccb..a52a8edc211 100644
--- a/src/librustc_mir/transform/mod.rs
+++ b/src/librustc_mir/transform/mod.rs
@@ -13,3 +13,4 @@ pub mod simplify_cfg;
 pub mod erase_regions;
 pub mod no_landing_pads;
 pub mod type_check;
+pub mod break_critical_edges;
diff --git a/src/librustc_trans/mir/block.rs b/src/librustc_trans/mir/block.rs
index a60b03dd724..1efbe5df162 100644
--- a/src/librustc_trans/mir/block.rs
+++ b/src/librustc_trans/mir/block.rs
@@ -300,33 +300,8 @@ impl<'bcx, 'tcx> MirContext<'bcx, 'tcx> {
 
                 // Many different ways to call a function handled here
                 if let Some(cleanup) = cleanup.map(|bb| self.bcx(bb)) {
-                    // We translate the copy into a temporary block. The temporary block is
-                    // necessary because the current block has already been terminated (by
-                    // `invoke`) and we cannot really translate into the target block
-                    // because:
-                    //  * The target block may have more than a single precedesor;
-                    //  * Some LLVM insns cannot have a preceeding store insn (phi,
-                    //    cleanuppad), and adding/prepending the store now may render
-                    //    those other instructions invalid.
-                    //
-                    // NB: This approach still may break some LLVM code. For example if the
-                    // target block starts with a `phi` (which may only match on immediate
-                    // precedesors), it cannot know about this temporary block thus
-                    // resulting in an invalid code:
-                    //
-                    // this:
-                    //     …
-                    //     %0 = …
-                    //     %1 = invoke to label %temp …
-                    // temp:
-                    //     store ty %1, ty* %dest
-                    //     br label %actualtargetblock
-                    // actualtargetblock:            ; preds: %temp, …
-                    //     phi … [%this, …], [%0, …] ; ERROR: phi requires to match only on
-                    //                               ; immediate precedesors
-
-                    let ret_bcx = if destination.is_some() {
-                        self.fcx.new_block("", None)
+                    let ret_bcx = if let Some((_, target)) = *destination {
+                        self.blocks[target.index()]
                     } else {
                         self.unreachable_block()
                     };
@@ -343,15 +318,16 @@ impl<'bcx, 'tcx> MirContext<'bcx, 'tcx> {
                         self.set_operand_dropped(bcx, op);
                     });
 
-                    if let Some((_, target)) = *destination {
+                    if destination.is_some() {
                         let ret_bcx = ret_bcx.build();
-                        if let Some(ret_dest) = ret_dest {
-                            fn_ty.ret.store(&ret_bcx, invokeret, ret_dest.llval);
-                        }
-                        for op in args {
-                            self.set_operand_dropped(&ret_bcx, op);
-                        }
-                        ret_bcx.br(self.llblock(target));
+                        ret_bcx.at_start(|ret_bcx| {
+                            if let Some(ret_dest) = ret_dest {
+                                fn_ty.ret.store(&ret_bcx, invokeret, ret_dest.llval);
+                            }
+                            for op in args {
+                                self.set_operand_dropped(&ret_bcx, op);
+                            }
+                        });
                     }
                 } else {
                     let llret = bcx.call(fn_ptr, &llargs, cleanup_bundle.as_ref());
diff --git a/src/librustc_trans/mir/mod.rs b/src/librustc_trans/mir/mod.rs
index 152834a07ba..27e09d04a9d 100644
--- a/src/librustc_trans/mir/mod.rs
+++ b/src/librustc_trans/mir/mod.rs
@@ -19,10 +19,11 @@ use common::{self, Block, BlockAndBuilder, FunctionContext};
 use std::ops::Deref;
 use std::rc::Rc;
 
+use rustc_data_structures::bitvec::BitVector;
+
 use self::lvalue::{LvalueRef, get_dataptr, get_meta};
 use rustc_mir::traversal;
 
-use self::lvalue::LvalueRef;
 use self::operand::OperandRef;
 
 #[derive(Clone)]
@@ -98,7 +99,7 @@ enum TempRef<'tcx> {
 
 ///////////////////////////////////////////////////////////////////////////
 
-pub fn trans_mir<'blk, 'tcx>(fcx: &'blk FunctionContext<'blk, 'tcx>) {
+pub fn trans_mir<'blk, 'tcx: 'blk>(fcx: &'blk FunctionContext<'blk, 'tcx>) {
     let bcx = fcx.init(false, None).build();
     let mir = bcx.mir();
 
@@ -135,8 +136,13 @@ pub fn trans_mir<'blk, 'tcx>(fcx: &'blk FunctionContext<'blk, 'tcx>) {
     let block_bcxs: Vec<Block<'blk,'tcx>> =
         mir_blocks.iter()
                   .map(|&bb|{
-                      // FIXME(#30941) this doesn't handle msvc-style exceptions
-                      fcx.new_block(&format!("{:?}", bb), None)
+                      if bb == mir::START_BLOCK {
+                          fcx.new_block("start", None)
+                      } else if bb == mir::END_BLOCK {
+                          fcx.new_block("end", None)
+                      } else {
+                          fcx.new_block(&format!("{:?}", bb), None)
+                      }
                   })
                   .collect();
 
@@ -145,7 +151,7 @@ pub fn trans_mir<'blk, 'tcx>(fcx: &'blk FunctionContext<'blk, 'tcx>) {
     bcx.br(start_bcx.llbb);
 
     let mut mircx = MirContext {
-        mir: mir,
+        mir: mir.clone(),
         fcx: fcx,
         llpersonalityslot: None,
         blocks: block_bcxs,
@@ -155,13 +161,28 @@ pub fn trans_mir<'blk, 'tcx>(fcx: &'blk FunctionContext<'blk, 'tcx>) {
         args: args,
     };
 
-    let rpo = traversal::reverse_postorder(mir);
+    let mut visited = BitVector::new(mir_blocks.len());
+
+    let rpo = traversal::reverse_postorder(&mir);
     // Translate the body of each block using reverse postorder
     for (bb, _) in rpo {
+        visited.insert(bb.index());
         mircx.trans_block(bb);
     }
 
+    // Add unreachable instructions at the end of unreachable blocks
+    // so they're actually terminated.
+    // TODO: Remove the blocks from the function
+    for &bb in &mir_blocks {
+        if !visited.contains(bb.index()) {
+            mircx.blocks[bb.index()].build().unreachable();
+        }
+    }
+
+
     fcx.cleanup();
+
+    debug!("trans_mir: {:?}", ::trans::value::Value(fcx.llfn));
 }
 
 /// Produce, for each argument, a `ValueRef` pointing at the