Only break critical edges where actually needed

Currently, to prepare for MIR trans, we break _all_ critical edges,
although we only actually need to do this for edges originating from a
call that gets translated to an invoke instruction in LLVM.

This has the unfortunate effect of undoing a bunch of the things that
SimplifyCfg has done. A particularly bad case arises when you have a
C-like enum with N variants and a derived PartialEq implementation.

In that case, the match on the (&lhs, &rhs) tuple gets translated into
nested matches with N arms each and a basic block each, resulting in N²
basic blocks. SimplifyCfg reduces that to roughly 2*N basic blocks, but
breaking the critical edges means that we go back to N².

In nickel.rs, there is such an enum with roughly N=800. So we get about
640K basic blocks or 2.5M lines of LLVM IR. LLVM takes a while to
reduce that to the final "disr_a == disr_b".

So before this patch, we had 2.5M lines of IR with 640K basic blocks,
which took about about 3.6s in LLVM to get optimized and translated.
After this patch, we get about 650K lines with about 1.6K basic blocks
and spent a little less than 0.2s in LLVM.

cc #33111

src/librustc_driver/driver.rs

@@ -1009,7 +1009,7 @@ pub fn phase_4_translate_to_llvm<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
         passes.push_pass(box mir::transform::no_landing_pads::NoLandingPads);
         passes.push_pass(box mir::transform::remove_dead_blocks::RemoveDeadBlocks);
         passes.push_pass(box mir::transform::erase_regions::EraseRegions);
-        passes.push_pass(box mir::transform::break_critical_edges::BreakCriticalEdges);
+        passes.push_pass(box mir::transform::break_cleanup_edges::BreakCleanupEdges);
         passes.run_passes(tcx, &mut mir_map);
     });
 

src/librustc_mir/transform/break_cleanup_edges.rs

@@ -0,0 +1,111 @@
+// 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::ty::TyCtxt;
+use rustc::mir::repr::*;
+use rustc::mir::transform::{MirPass, MirSource, Pass};
+
+use rustc_data_structures::bitvec::BitVector;
+
+use pretty;
+
+use traversal;
+
+pub struct BreakCleanupEdges;
+
+/**
+ * Breaks outgoing critical edges for call terminators 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. We need this for calls that are
+ * translated to LLVM invoke instructions, because invoke is a block terminator
+ * in LLVM so we can't insert any code to handle the call's result into the
+ * block that performs the call.
+ *
+ * This function will break those edges by inserting new blocks along them.
+ *
+ * NOTE: Simplify CFG will happily undo most of the work this pass does.
+ *
+ */
+
+impl<'tcx> MirPass<'tcx> for BreakCleanupEdges {
+    fn run_pass<'a>(&mut self, tcx: TyCtxt<'a, 'tcx, 'tcx>, src: MirSource, mir: &mut Mir<'tcx>) {
+        let mut pred_count = vec![0u32; mir.basic_blocks.len()];
+
+        // Build the precedecessor map for the MIR
+        for (_, data) in traversal::preorder(mir) {
+            if let Some(ref term) = data.terminator {
+                for &tgt in term.successors().iter() {
+                    pred_count[tgt.index()] += 1;
+                }
+            }
+        }
+
+        let cleanup_map : BitVector = mir.basic_blocks
+            .iter().map(|bb| bb.is_cleanup).collect();
+
+        // 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 {
+                if term_is_invoke(term) {
+                    let term_span = term.span;
+                    let term_scope = term.scope;
+                    let succs = term.successors_mut();
+                    for tgt in succs {
+                        let num_preds = pred_count[tgt.index()];
+                        if num_preds > 1 {
+                            // It's a critical edge, break it
+                            let goto = Terminator {
+                                span: term_span,
+                                scope: term_scope,
+                                kind: TerminatorKind::Goto { target: *tgt }
+                            };
+                            let mut data = BasicBlockData::new(Some(goto));
+                            data.is_cleanup = cleanup_map.contains(tgt.index());
+
+                            // 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);
+                        }
+                    }
+                }
+            }
+        }
+
+        pretty::dump_mir(tcx, "break_cleanup_edges", &0, src, mir, None);
+        debug!("Broke {} N edges", new_blocks.len());
+
+        mir.basic_blocks.extend_from_slice(&new_blocks);
+    }
+}
+
+impl Pass for BreakCleanupEdges {}
+
+// Returns true if the terminator is a call that would use an invoke in LLVM.
+fn term_is_invoke(term: &Terminator) -> bool {
+    match term.kind {
+        TerminatorKind::Call { cleanup: Some(_), .. } |
+        TerminatorKind::Drop { unwind: Some(_), .. } => true,
+        _ => false
+    }
+}

src/librustc_mir/transform/mod.rs

@@ -13,6 +13,6 @@ pub mod simplify_cfg;
 pub mod erase_regions;
 pub mod no_landing_pads;
 pub mod type_check;
-pub mod break_critical_edges;
+pub mod break_cleanup_edges;
 pub mod promote_consts;
 pub mod qualify_consts;