break critical edges only when needed

arielb1 · arielb1 · commit 4106ab24d744 · 2016-06-05T09:27:26.000+03:00
the *only* place where critical edges need to be broken is on Call
instructions, so only break them there.
diff --git a/src/librustc_borrowck/borrowck/mir/elaborate_drops.rs b/src/librustc_borrowck/borrowck/mir/elaborate_drops.rs
@@ -327,12 +327,30 @@ impl<'b, 'tcx> ElaborateDropsCtxt<'b, 'tcx> {
         let data = self.mir.basic_block_data(bb);
         let terminator = data.terminator();
 
-        let unwind = Some(unwind.unwrap_or_else(|| {
-            // we can't use the resume block directly, because we
-            // may want to add a drop flag write.
-            self.jump_to_resume_block(terminator.scope,
-                                      terminator.span)
-        }));
+        let assign = Statement {
+            kind: StatementKind::Assign(location.clone(), Rvalue::Use(value.clone())),
+            span: terminator.span,
+            scope: terminator.scope
+        };
+
+        let unwind = unwind.unwrap_or(self.patch.resume_block());
+        let unwind = self.patch.new_block(BasicBlockData {
+            statements: vec![assign.clone()],
+            terminator: Some(Terminator {
+                kind: TerminatorKind::Goto { target: unwind },
+                ..*terminator
+            }),
+            is_cleanup: true
+        });
+
+        let target = self.patch.new_block(BasicBlockData {
+            statements: vec![assign],
+            terminator: Some(Terminator {
+                kind: TerminatorKind::Goto { target: target },
+                ..*terminator
+            }),
+            is_cleanup: data.is_cleanup,
+        });
 
         if !self.lvalue_is_tracked(location) {
             // drop and replace behind a pointer/array/whatever. The location
@@ -341,7 +359,7 @@ impl<'b, 'tcx> ElaborateDropsCtxt<'b, 'tcx> {
             self.patch.patch_terminator(bb, TerminatorKind::Drop {
                 location: location.clone(),
                 target: target,
-                unwind: unwind
+                unwind: Some(unwind)
             });
         } else {
             debug!("elaborate_drop_and_replace({:?}) - tracked", terminator);
@@ -356,24 +374,15 @@ impl<'b, 'tcx> ElaborateDropsCtxt<'b, 'tcx> {
                 lvalue: location,
                 path: path,
                 succ: target,
-                unwind: unwind
+                unwind: Some(unwind)
             }, bb);
             on_all_children_bits(self.tcx, self.mir, self.move_data(), path, |child| {
                 self.set_drop_flag(Location { block: target, index: 0 },
                                    child, DropFlagState::Present);
-                if let Some(unwind) = unwind {
-                    self.set_drop_flag(Location { block: unwind, index: 0 },
-                                       child, DropFlagState::Present);
-                }
+                self.set_drop_flag(Location { block: unwind, index: 0 },
+                                   child, DropFlagState::Present);
             });
         }
-
-        self.patch.add_assign(Location { block: target, index: 0 },
-                              location.clone(), Rvalue::Use(value.clone()));
-        if let Some(unwind) = unwind {
-            self.patch.add_assign(Location { block: unwind, index: 0 },
-                                  location.clone(), Rvalue::Use(value.clone()));
-        }
     }
 
     /// This elaborates a single drop instruction, located at `bb`, and
@@ -828,19 +837,6 @@ impl<'b, 'tcx> ElaborateDropsCtxt<'b, 'tcx> {
         })
     }
 
-    fn jump_to_resume_block<'a>(&mut self, scope: ScopeId, span: Span) -> BasicBlock {
-        let resume_block = self.patch.resume_block();
-        self.patch.new_block(BasicBlockData {
-            statements: vec![],
-            terminator: Some(Terminator {
-                scope: scope, span: span, kind: TerminatorKind::Goto {
-                    target: resume_block
-                }
-            }),
-            is_cleanup: true
-       })
-    }
-
     fn box_free_block<'a>(
         &mut self,
         c: &DropCtxt<'a, 'tcx>,
diff --git a/src/librustc_driver/driver.rs b/src/librustc_driver/driver.rs
@@ -1032,11 +1032,12 @@ 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_cleanup_edges::BreakCleanupEdges);
+        passes.push_pass(box mir::transform::add_call_guards::AddCallGuards);
         passes.push_pass(box borrowck::ElaborateDrops);
         passes.push_pass(box mir::transform::no_landing_pads::NoLandingPads);
         passes.push_pass(box mir::transform::simplify_cfg::SimplifyCfg);
-        passes.push_pass(box mir::transform::break_cleanup_edges::BreakCleanupEdges);
+        passes.push_pass(box mir::transform::add_call_guards::AddCallGuards);
+        passes.push_pass(box mir::transform::dump_mir::DumpMir("pre_trans"));
         passes.run_passes(tcx, &mut mir_map);
     });
 
diff --git a/src/librustc_mir/transform/add_call_guards.rs b/src/librustc_mir/transform/add_call_guards.rs
@@ -12,13 +12,11 @@ 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;
+pub struct AddCallGuards;
 
 /**
  * Breaks outgoing critical edges for call terminators in the MIR.
@@ -40,7 +38,7 @@ pub struct BreakCleanupEdges;
  *
  */
 
-impl<'tcx> MirPass<'tcx> for BreakCleanupEdges {
+impl<'tcx> MirPass<'tcx> for AddCallGuards {
     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()];
 
@@ -53,9 +51,6 @@ impl<'tcx> MirPass<'tcx> for BreakCleanupEdges {
             }
         }
 
-        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();
 
@@ -65,30 +60,31 @@ impl<'tcx> MirPass<'tcx> for BreakCleanupEdges {
         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);
-                        }
-                    }
+            match data.terminator {
+                Some(Terminator {
+                    kind: TerminatorKind::Call {
+                        destination: Some((_, ref mut destination)),
+                        cleanup: Some(_),
+                        ..
+                    }, span, scope
+                }) if pred_count[destination.index()] > 1 => {
+                    // It's a critical edge, break it
+                    let call_guard = BasicBlockData {
+                        statements: vec![],
+                        is_cleanup: data.is_cleanup,
+                        terminator: Some(Terminator {
+                            span: span,
+                            scope: scope,
+                            kind: TerminatorKind::Goto { target: *destination }
+                        })
+                    };
+
+                    // Get the index it will be when inserted into the MIR
+                    let idx = cur_len + new_blocks.len();
+                    new_blocks.push(call_guard);
+                    *destination = BasicBlock::new(idx);
                 }
+                _ => {}
             }
         }
 
@@ -99,15 +95,4 @@ impl<'tcx> MirPass<'tcx> for BreakCleanupEdges {
     }
 }
 
-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(_), .. } |
-        // FIXME: not sure whether we need this one
-        TerminatorKind::Drop { unwind: Some(_), .. } |
-        TerminatorKind::DropAndReplace { .. } => true,
-        _ => false
-    }
-}
+impl Pass for AddCallGuards {}
diff --git a/src/librustc_mir/transform/dump_mir.rs b/src/librustc_mir/transform/dump_mir.rs
@@ -0,0 +1,27 @@
+// Copyright 2015 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.
+
+//! This pass just dumps MIR at a specified point.
+
+use rustc::ty::TyCtxt;
+use rustc::mir::repr::*;
+use rustc::mir::transform::{Pass, MirPass, MirSource};
+use pretty;
+
+pub struct DumpMir<'a>(pub &'a str);
+
+impl<'b, 'tcx> MirPass<'tcx> for DumpMir<'b> {
+    fn run_pass<'a>(&mut self, tcx: TyCtxt<'a, 'tcx, 'tcx>,
+                    src: MirSource, mir: &mut Mir<'tcx>) {
+        pretty::dump_mir(tcx, self.0, &0, src, mir, None);
+    }
+}
+
+impl<'b> Pass for DumpMir<'b> {}
diff --git a/src/librustc_mir/transform/mod.rs b/src/librustc_mir/transform/mod.rs
@@ -13,6 +13,7 @@ pub mod simplify_cfg;
 pub mod erase_regions;
 pub mod no_landing_pads;
 pub mod type_check;
-pub mod break_cleanup_edges;
+pub mod add_call_guards;
 pub mod promote_consts;
 pub mod qualify_consts;
+pub mod dump_mir;
