Subpart9 for async drop (major3) - elaborate_drops changes

Author: azhogin

Cargo.lock

@@ -4099,6 +4099,7 @@ dependencies = [
 name = "rustc_mir_dataflow"
 version = "0.0.0"
 dependencies = [
+ "itertools",
  "polonius-engine",
  "regex",
  "rustc_ast",

compiler/rustc_mir_dataflow/Cargo.toml

@@ -5,6 +5,7 @@ edition = "2021"
 
 [dependencies]
 # tidy-alphabetical-start
+itertools = "0.12"
 polonius-engine = "0.13.0"
 regex = "1"
 rustc_ast = { path = "../rustc_ast" }

compiler/rustc_mir_dataflow/src/elaborate_drops.rs

@@ -1174,7 +1174,13 @@ fn elaborate_coroutine_drops<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     let def_id = body.source.def_id();
     let param_env = tcx.param_env(def_id);
 
-    let mut elaborator = DropShimElaborator { body, patch: MirPatch::new(body), tcx, param_env };
+    let mut elaborator = DropShimElaborator {
+        body,
+        patch: MirPatch::new(body),
+        tcx,
+        param_env,
+        produce_async_drops: false,
+    };
 
     for (block, block_data) in body.basic_blocks.iter_enumerated() {
         let (target, unwind, source_info) = match block_data.terminator() {
@@ -1213,6 +1219,7 @@ fn elaborate_coroutine_drops<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             *target,
             unwind,
             block,
+            None,
         );
     }
     elaborator.patch.apply(body);

compiler/rustc_mir_transform/src/coroutine.rs

@@ -165,6 +165,14 @@ impl<'a, 'tcx> DropElaborator<'a, 'tcx> for Elaborator<'a, '_, '_, 'tcx> {
         self.ctxt.param_env()
     }
 
+    fn allow_async_drops(&self) -> bool {
+        true
+    }
+
+    fn terminator_loc(&self, bb: BasicBlock) -> Location {
+        self.ctxt.patch.terminator_loc(self.ctxt.body, bb)
+    }
+
     #[instrument(level = "debug", skip(self), ret)]
     fn drop_style(&self, path: Self::Path, mode: DropFlagMode) -> DropStyle {
         let ((maybe_live, maybe_dead), multipart) = match mode {
@@ -333,7 +341,7 @@ impl<'b, 'mir, 'tcx> ElaborateDropsCtxt<'b, 'mir, 'tcx> {
         // This function should mirror what `collect_drop_flags` does.
         for (bb, data) in self.body.basic_blocks.iter_enumerated() {
             let terminator = data.terminator();
-            let TerminatorKind::Drop { place, target, unwind, replace, drop: _, async_fut: _ } =
+            let TerminatorKind::Drop { place, target, unwind, replace, drop, async_fut: _ } =
                 terminator.kind
             else {
                 continue;
@@ -379,6 +387,7 @@ impl<'b, 'mir, 'tcx> ElaborateDropsCtxt<'b, 'mir, 'tcx> {
                         target,
                         unwind,
                         bb,
+                        drop,
                     )
                 }
                 LookupResult::Parent(None) => {}

compiler/rustc_mir_transform/src/elaborate_drops.rs

@@ -281,8 +281,13 @@ fn build_drop_shim<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, ty: Option<Ty<'tcx>>)
     if ty.is_some() {
         let patch = {
             let param_env = tcx.param_env_reveal_all_normalized(def_id);
-            let mut elaborator =
-                DropShimElaborator { body: &body, patch: MirPatch::new(&body), tcx, param_env };
+            let mut elaborator = DropShimElaborator {
+                body: &body,
+                patch: MirPatch::new(&body),
+                tcx,
+                param_env,
+                produce_async_drops: false,
+            };
             let dropee = tcx.mk_place_deref(dropee_ptr);
             let resume_block = elaborator.patch.resume_block();
             elaborate_drops::elaborate_drop(
@@ -293,6 +298,7 @@ fn build_drop_shim<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, ty: Option<Ty<'tcx>>)
                 return_block,
                 elaborate_drops::Unwind::To(resume_block),
                 START_BLOCK,
+                None,
             );
             elaborator.patch
         };
@@ -341,6 +347,7 @@ pub struct DropShimElaborator<'a, 'tcx> {
     pub patch: MirPatch<'tcx>,
     pub tcx: TyCtxt<'tcx>,
     pub param_env: ty::ParamEnv<'tcx>,
+    pub produce_async_drops: bool,
 }
 
 impl fmt::Debug for DropShimElaborator<'_, '_> {
@@ -365,6 +372,13 @@ impl<'a, 'tcx> DropElaborator<'a, 'tcx> for DropShimElaborator<'a, 'tcx> {
         self.param_env
     }
 
+    fn terminator_loc(&self, bb: BasicBlock) -> Location {
+        self.patch.terminator_loc(self.body, bb)
+    }
+    fn allow_async_drops(&self) -> bool {
+        self.produce_async_drops
+    }
+
     fn drop_style(&self, _path: Self::Path, mode: DropFlagMode) -> DropStyle {
         match mode {
             DropFlagMode::Shallow => {