Auto merge of #43748 - RalfJung:mir-validate2, r=arielb1

AddValidation: handle Call terminators into blocks that have multiple incoming edges

The old code was just wrong: It would add validation on paths that don't even come from the call, and it would add multiple validations if multiple calls end return to the same block.

src/librustc_driver/driver.rs

@@ -949,15 +949,16 @@ pub fn phase_3_run_analysis_passes<'tcx, F, R>(sess: &'tcx Session,
 
     // These next passes must be executed together
     passes.push_pass(MIR_OPTIMIZED, mir::transform::no_landing_pads::NoLandingPads);
-    passes.push_pass(MIR_OPTIMIZED, mir::transform::add_call_guards::AddCallGuards);
+    passes.push_pass(MIR_OPTIMIZED, mir::transform::add_call_guards::CriticalCallEdges);
     passes.push_pass(MIR_OPTIMIZED, mir::transform::elaborate_drops::ElaborateDrops);
     passes.push_pass(MIR_OPTIMIZED, mir::transform::no_landing_pads::NoLandingPads);
+    // AddValidation needs to run after ElaborateDrops and before EraseRegions, and it needs
+    // an AllCallEdges pass right before it.
+    passes.push_pass(MIR_OPTIMIZED, mir::transform::add_call_guards::AllCallEdges);
+    passes.push_pass(MIR_OPTIMIZED, mir::transform::add_validation::AddValidation);
     passes.push_pass(MIR_OPTIMIZED, mir::transform::simplify::SimplifyCfg::new("elaborate-drops"));
     // No lifetime analysis based on borrowing can be done from here on out.
 
-    // AddValidation needs to run after ElaborateDrops and before EraseRegions.
-    passes.push_pass(MIR_OPTIMIZED, mir::transform::add_validation::AddValidation);
-
     // From here on out, regions are gone.
     passes.push_pass(MIR_OPTIMIZED, mir::transform::erase_regions::EraseRegions);
 
@@ -967,7 +968,7 @@ pub fn phase_3_run_analysis_passes<'tcx, F, R>(sess: &'tcx Session,
     passes.push_pass(MIR_OPTIMIZED, mir::transform::deaggregator::Deaggregator);
     passes.push_pass(MIR_OPTIMIZED, mir::transform::copy_prop::CopyPropagation);
     passes.push_pass(MIR_OPTIMIZED, mir::transform::simplify::SimplifyLocals);
-    passes.push_pass(MIR_OPTIMIZED, mir::transform::add_call_guards::AddCallGuards);
+    passes.push_pass(MIR_OPTIMIZED, mir::transform::add_call_guards::CriticalCallEdges);
     passes.push_pass(MIR_OPTIMIZED, mir::transform::dump_mir::Marker("PreTrans"));
 
     TyCtxt::create_and_enter(sess,

src/librustc_mir/shim.rs

@@ -105,7 +105,7 @@ fn make_shim<'a, 'tcx>(tcx: ty::TyCtxt<'a, 'tcx, 'tcx>,
         debug!("make_shim({:?}) = untransformed {:?}", instance, result);
         no_landing_pads::no_landing_pads(tcx, &mut result);
         simplify::simplify_cfg(&mut result);
-        add_call_guards::add_call_guards(&mut result);
+        add_call_guards::CriticalCallEdges.add_call_guards(&mut result);
     debug!("make_shim({:?}) = {:?}", instance, result);
 
     tcx.alloc_mir(result)

src/librustc_mir/transform/add_call_guards.rs

@@ -13,7 +13,12 @@ use rustc::mir::*;
 use rustc::mir::transform::{MirPass, MirSource};
 use rustc_data_structures::indexed_vec::{Idx, IndexVec};
 
-pub struct AddCallGuards;
+#[derive(PartialEq)]
+pub enum AddCallGuards {
+    AllCallEdges,
+    CriticalCallEdges,
+}
+pub use self::AddCallGuards::*;
 
 /**
  * Breaks outgoing critical edges for call terminators in the MIR.
@@ -40,48 +45,52 @@ impl MirPass for AddCallGuards {
                           _tcx: TyCtxt<'a, 'tcx, 'tcx>,
                           _src: MirSource,
                           mir: &mut Mir<'tcx>) {
-        add_call_guards(mir);
+        self.add_call_guards(mir);
     }
 }
 
-pub fn add_call_guards(mir: &mut Mir) {
-    let pred_count: IndexVec<_, _> =
-        mir.predecessors().iter().map(|ps| ps.len()).collect();
+impl AddCallGuards {
+    pub fn add_call_guards(&self, mir: &mut Mir) {
+        let pred_count: IndexVec<_, _> =
+            mir.predecessors().iter().map(|ps| ps.len()).collect();
 
-    // We need a place to store the new blocks generated
-    let mut new_blocks = Vec::new();
+        // We need a place to store the new blocks generated
+        let mut new_blocks = Vec::new();
 
-    let cur_len = mir.basic_blocks().len();
+        let cur_len = mir.basic_blocks().len();
 
-    for block in mir.basic_blocks_mut() {
-        match block.terminator {
-            Some(Terminator {
-                kind: TerminatorKind::Call {
-                    destination: Some((_, ref mut destination)),
-                    cleanup: Some(_),
-                    ..
-                }, source_info
-            }) if pred_count[*destination] > 1 => {
-                // It's a critical edge, break it
-                let call_guard = BasicBlockData {
-                    statements: vec![],
-                    is_cleanup: block.is_cleanup,
-                    terminator: Some(Terminator {
-                        source_info: source_info,
-                        kind: TerminatorKind::Goto { target: *destination }
-                    })
-                };
+        for block in mir.basic_blocks_mut() {
+            match block.terminator {
+                Some(Terminator {
+                    kind: TerminatorKind::Call {
+                        destination: Some((_, ref mut destination)),
+                        cleanup,
+                        ..
+                    }, source_info
+                }) if pred_count[*destination] > 1 &&
+                      (cleanup.is_some() || self == &AllCallEdges) =>
+                {
+                    // It's a critical edge, break it
+                    let call_guard = BasicBlockData {
+                        statements: vec![],
+                        is_cleanup: block.is_cleanup,
+                        terminator: Some(Terminator {
+                            source_info: source_info,
+                            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);
+                    // 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);
+                }
+                _ => {}
             }
-            _ => {}
         }
-    }
 
-    debug!("Broke {} N edges", new_blocks.len());
+        debug!("Broke {} N edges", new_blocks.len());
 
-    mir.basic_blocks_mut().extend(new_blocks);
+        mir.basic_blocks_mut().extend(new_blocks);
+    }
 }