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_critical_edges.rs → src/librustc_mir/transform/break_cleanup_edges.rs

@@ -14,12 +14,14 @@
 
 use rustc_data_structures::bitvec::BitVector;
 
+use pretty;
+
 use traversal;
 
-pub struct BreakCriticalEdges;
+pub struct BreakCleanupEdges;
 
 /**
- * Breaks critical edges in the MIR.
+ * 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.
@@ -27,87 +29,79 @@
  * 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.
+ * "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.
  *
- * 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.
  *
  */
 
-impl<'tcx> MirPass<'tcx> for BreakCriticalEdges {
-    fn run_pass<'a>(&mut self, _: TyCtxt<'a, 'tcx, 'tcx>,
-                    _: MirSource, mir: &mut Mir<'tcx>) {
-        break_critical_edges(mir);
-    }
-}
-
-impl Pass for BreakCriticalEdges {}
-
-fn break_critical_edges(mir: &mut Mir) {
-    let mut pred_count = vec![0u32; mir.basic_blocks.len()];
+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;
+        // 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 {
-            let is_invoke = term_is_invoke(term);
-            let term_span = term.span;
-            let term_scope = term.scope;
-            let succs = term.successors_mut();
-            if succs.len() > 1 || (succs.len() > 0 && is_invoke) {
-                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);
+        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);
+                        }
                     }
                 }
             }
         }
-    }
 
-    debug!("Broke {} N edges", new_blocks.len());
+        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);
+        mir.basic_blocks.extend_from_slice(&new_blocks);
+    }
 }
 
-// Returns true if the terminator would use an invoke in LLVM.
+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(_), .. } |

src/librustc_mir/transform/mod.rs

@@ -13,6 +13,6 @@
 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;