From a43e486e8d126f8a815c4e8ebba5ce8692864b06 Mon Sep 17 00:00:00 2001
From: Dylan MacKenzie <ecstaticmorse@gmail.com>
Date: Sat, 2 May 2020 21:16:17 -0700
Subject: [PATCH] Add MIR phase and query for drop elaboration

---
 src/librustc_middle/mir/mod.rs    |  3 +-
 src/librustc_middle/query/mod.rs  |  6 ++++
 src/librustc_mir/transform/mod.rs | 49 +++++++++++++++++++++++++------
 3 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/src/librustc_middle/mir/mod.rs b/src/librustc_middle/mir/mod.rs
index 129f9691ea5..21f5d9e7dd4 100644
--- a/src/librustc_middle/mir/mod.rs
+++ b/src/librustc_middle/mir/mod.rs
@@ -76,7 +76,8 @@ pub enum MirPhase {
     Build = 0,
     Const = 1,
     Validated = 2,
-    Optimized = 3,
+    DropElab = 3,
+    Optimized = 4,
 }
 
 impl MirPhase {
diff --git a/src/librustc_middle/query/mod.rs b/src/librustc_middle/query/mod.rs
index 4d7e7882e42..9e3881b28b6 100644
--- a/src/librustc_middle/query/mod.rs
+++ b/src/librustc_middle/query/mod.rs
@@ -190,6 +190,12 @@ fn describe_as_module(def_id: DefId, tcx: TyCtxt<'_>) -> String {
             no_hash
         }
 
+        query mir_drops_elaborated_and_const_checked(key: LocalDefId) -> Steal<mir::Body<'tcx>> {
+            storage(ArenaCacheSelector<'tcx>)
+            no_hash
+            desc { |tcx| "elaborating drops for `{}`", tcx.def_path_str(key.to_def_id()) }
+        }
+
         query mir_validated(key: LocalDefId) ->
             (
                 Steal<mir::Body<'tcx>>,
diff --git a/src/librustc_mir/transform/mod.rs b/src/librustc_mir/transform/mod.rs
index 26725a2ac02..4240b528a61 100644
--- a/src/librustc_mir/transform/mod.rs
+++ b/src/librustc_mir/transform/mod.rs
@@ -49,6 +49,7 @@ pub(crate) fn provide(providers: &mut Providers<'_>) {
         mir_const,
         mir_const_qualif,
         mir_validated,
+        mir_drops_elaborated_and_const_checked,
         optimized_mir,
         is_mir_available,
         promoted_mir,
@@ -294,12 +295,31 @@ fn mir_validated(
     (tcx.alloc_steal_mir(body), tcx.alloc_steal_promoted(promoted))
 }
 
-fn run_optimization_passes<'tcx>(
+fn mir_drops_elaborated_and_const_checked<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    def_id: LocalDefId,
+) -> Steal<Body<'tcx>> {
+    // (Mir-)Borrowck uses `mir_validated`, so we have to force it to
+    // execute before we can steal.
+    tcx.ensure().mir_borrowck(def_id);
+
+    let (body, _) = tcx.mir_validated(def_id);
+    let mut body = body.steal();
+
+    run_post_borrowck_cleanup_passes(tcx, &mut body, def_id, None);
+    check_consts::post_drop_elaboration::check_live_drops(tcx, def_id, &body);
+    tcx.alloc_steal_mir(body)
+}
+
+/// After this series of passes, no lifetime analysis based on borrowing can be done.
+fn run_post_borrowck_cleanup_passes<'tcx>(
     tcx: TyCtxt<'tcx>,
     body: &mut Body<'tcx>,
     def_id: LocalDefId,
     promoted: Option<Promoted>,
 ) {
+    debug!("post_borrowck_cleanup({:?})", def_id);
+
     let post_borrowck_cleanup: &[&dyn MirPass<'tcx>] = &[
         // Remove all things only needed by analysis
         &no_landing_pads::NoLandingPads::new(tcx),
@@ -318,9 +338,24 @@ fn run_optimization_passes<'tcx>(
         // but before optimizations begin.
         &add_retag::AddRetag,
         &simplify::SimplifyCfg::new("elaborate-drops"),
-        // No lifetime analysis based on borrowing can be done from here on out.
     ];
 
+    run_passes(
+        tcx,
+        body,
+        InstanceDef::Item(def_id.to_def_id()),
+        promoted,
+        MirPhase::DropElab,
+        &[post_borrowck_cleanup],
+    );
+}
+
+fn run_optimization_passes<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    body: &mut Body<'tcx>,
+    def_id: LocalDefId,
+    promoted: Option<Promoted>,
+) {
     let optimizations: &[&dyn MirPass<'tcx>] = &[
         &unreachable_prop::UnreachablePropagation,
         &uninhabited_enum_branching::UninhabitedEnumBranching,
@@ -368,6 +403,7 @@ fn run_optimization_passes<'tcx>(
 
     let mir_opt_level = tcx.sess.opts.debugging_opts.mir_opt_level;
 
+    #[rustfmt::skip]
     run_passes(
         tcx,
         body,
@@ -375,7 +411,6 @@ fn run_optimization_passes<'tcx>(
         promoted,
         MirPhase::Optimized,
         &[
-            post_borrowck_cleanup,
             if mir_opt_level > 0 { optimizations } else { no_optimizations },
             pre_codegen_cleanup,
         ],
@@ -393,12 +428,7 @@ fn optimized_mir(tcx: TyCtxt<'_>, def_id: DefId) -> Body<'_> {
 
     let def_id = def_id.expect_local();
 
-    // (Mir-)Borrowck uses `mir_validated`, so we have to force it to
-    // execute before we can steal.
-    tcx.ensure().mir_borrowck(def_id);
-
-    let (body, _) = tcx.mir_validated(def_id);
-    let mut body = body.steal();
+    let mut body = tcx.mir_drops_elaborated_and_const_checked(def_id).steal();
     run_optimization_passes(tcx, &mut body, def_id, None);
 
     debug_assert!(!body.has_free_regions(), "Free regions in optimized MIR");
@@ -418,6 +448,7 @@ fn promoted_mir(tcx: TyCtxt<'_>, def_id: DefId) -> IndexVec<Promoted, Body<'_>>
     let mut promoted = promoted.steal();
 
     for (p, mut body) in promoted.iter_enumerated_mut() {
+        run_post_borrowck_cleanup_passes(tcx, &mut body, def_id, Some(p));
         run_optimization_passes(tcx, &mut body, def_id, Some(p));
     }
 
-- 
GitLab