Expand async drops during drop elaboration.

Author: cjgillot

compiler/rustc_hir/src/hir.rs

@@ -2327,6 +2327,13 @@ impl CoroutineKind {
         matches!(self, CoroutineKind::Desugared(_, CoroutineSource::Fn))
     }
 
+    pub fn is_async_desugaring(self) -> bool {
+        matches!(
+            self,
+            CoroutineKind::Desugared(CoroutineDesugaring::Async | CoroutineDesugaring::AsyncGen, _)
+        )
+    }
+
     pub fn to_plural_string(&self) -> String {
         match self {
             CoroutineKind::Desugared(d, CoroutineSource::Fn) => format!("{d:#}fn bodies"),

compiler/rustc_middle/src/mir/statement.rs

@@ -624,11 +624,7 @@ impl<'tcx> Operand<'tcx> {
         span: Span,
     ) -> Self {
         let ty = Ty::new_fn_def(tcx, def_id, args);
-        Operand::Constant(Box::new(ConstOperand {
-            span,
-            user_ty: None,
-            const_: Const::Val(ConstValue::ZeroSized, ty),
-        }))
+        Operand::zero_sized_constant(ty, span)
     }
 
     /// Convenience helper to make a constant that refers to the given `DefId` and args. Since this
@@ -643,6 +639,12 @@ impl<'tcx> Operand<'tcx> {
         Operand::Constant(Box::new(ConstOperand { span, user_ty: None, const_ }))
     }
 
+    /// Convenience helper to make a constant that refers to a zero-sized type.
+    pub fn zero_sized_constant(ty: Ty<'tcx>, span: Span) -> Self {
+        let const_ = Const::Val(ConstValue::ZeroSized, ty);
+        Operand::Constant(Box::new(ConstOperand { span, user_ty: None, const_ }))
+    }
+
     pub fn is_move(&self) -> bool {
         matches!(self, Operand::Move(..))
     }

compiler/rustc_middle/src/mir/syntax.rs

@@ -760,21 +760,16 @@ pub enum TerminatorKind<'tcx> {
     /// meaning.
     ///
     /// Async drop processing:
-    /// In compiler/rustc_mir_build/src/build/scope.rs we detect possible async drop:
-    ///   drop of object with `needs_async_drop`.
-    /// Async drop later, in StateTransform pass, may be expanded into additional yield-point
-    ///   for poll-loop of async drop future.
-    /// So we need prepared 'drop' target block in the similar way as for `Yield` terminator
-    ///   (see `drops.build_mir::<CoroutineDrop>` in scopes.rs).
-    /// In compiler/rustc_mir_transform/src/elaborate_drops.rs for object implementing `AsyncDrop` trait
-    ///   we need to prepare async drop feature - resolve `AsyncDrop::drop` and codegen call.
-    /// `async_fut` is set to the corresponding local.
-    /// For coroutine drop we don't need this logic because coroutine drop works with the same
-    ///   layout object as coroutine itself. So `async_fut` will be `None` for coroutine drop.
-    /// Both `drop` and `async_fut` fields are only used in compiler/rustc_mir_transform/src/coroutine.rs,
-    ///   StateTransform pass. In `expand_async_drops` async drops are expanded
-    ///   into one or two yield points with poll ready/pending switch.
-    /// When a coroutine has any internal async drop, the coroutine drop function will be async
+    ///   MIR building detects possible async drops, and constructs a complete CFG. To correctly
+    ///   handle the coroutine being dropped while itself drops, we need a 'drop' target
+    ///   similar to `Yield` terminator (see `drops.build_mir::<CoroutineDrop>`).
+    ///
+    ///   Drop elaboration later refines the set of useful async drops. If there is no need for an
+    ///   async drop, it is downgraded to a sync drop by setting `drop` to `None` If this is an
+    ///   actual async drop, it is expanded to an `await` loop over the `async_drop_in_place` or
+    ///   `AsyncDrop::drop` coroutine.
+    ///
+    ///   When a coroutine has any internal async drop, the coroutine drop function will be async
     ///   (generated by `create_coroutine_drop_shim_async`, not `create_coroutine_drop_shim`).
     Drop {
         place: Place<'tcx>,

compiler/rustc_mir_transform/src/coroutine.rs

@@ -56,9 +56,9 @@ use std::ops;
 
 pub(super) use by_move_body::coroutine_by_move_body_def_id;
 use drop::{
-    cleanup_async_drops, create_coroutine_drop_shim, create_coroutine_drop_shim_async,
-    create_coroutine_drop_shim_proxy_async, elaborate_coroutine_drops, expand_async_drops,
-    has_expandable_async_drops, insert_clean_drop,
+    create_coroutine_drop_shim, create_coroutine_drop_shim_async,
+    create_coroutine_drop_shim_proxy_async, elaborate_coroutine_drops, has_async_drops,
+    insert_clean_drop,
 };
 use itertools::izip;
 use rustc_abi::{FieldIdx, VariantIdx};
@@ -70,7 +70,6 @@ use rustc_index::bit_set::{BitMatrix, DenseBitSet, GrowableBitSet};
 use rustc_index::{Idx, IndexVec, indexvec};
 use rustc_middle::mir::visit::{MutVisitor, MutatingUseContext, PlaceContext, Visitor};
 use rustc_middle::mir::*;
-use rustc_middle::ty::util::Discr;
 use rustc_middle::ty::{
     self, CoroutineArgs, CoroutineArgsExt, GenericArgsRef, InstanceKind, Ty, TyCtxt, TypingMode,
 };
@@ -82,8 +81,8 @@ use rustc_mir_dataflow::impls::{
 use rustc_mir_dataflow::{
     Analysis, Results, ResultsCursor, ResultsVisitor, visit_reachable_results,
 };
+use rustc_span::Span;
 use rustc_span::def_id::{DefId, LocalDefId};
-use rustc_span::{DUMMY_SP, Span, dummy_spanned};
 use rustc_trait_selection::error_reporting::InferCtxtErrorExt;
 use rustc_trait_selection::infer::TyCtxtInferExt as _;
 use rustc_trait_selection::traits::{ObligationCause, ObligationCauseCode, ObligationCtxt};
@@ -169,7 +168,7 @@ fn replace_base<'tcx>(place: &mut Place<'tcx>, new_base: Place<'tcx>, tcx: TyCtx
 }
 
 const SELF_ARG: Local = Local::arg(0);
-const CTX_ARG: Local = Local::arg(1);
+pub(crate) const CTX_ARG: Local = Local::arg(1);
 
 /// A `yield` point in the coroutine.
 struct SuspensionPoint<'tcx> {
@@ -585,7 +584,7 @@ fn make_coroutine_state_argument_pinned<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body
 /// still using the `ResumeTy` indirection for the time being, and that indirection
 /// is removed here. After this transform, the coroutine body only knows about `&mut Context<'_>`.
 #[tracing::instrument(level = "trace", skip(tcx, body), ret)]
-fn transform_async_context<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) -> Ty<'tcx> {
+fn transform_async_context<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     let context_mut_ref = Ty::new_task_context(tcx);
 
     // replace the type of the `resume` argument
@@ -615,7 +614,6 @@ fn transform_async_context<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) -> Ty
             _ => {}
         }
     }
-    context_mut_ref
 }
 
 fn eliminate_get_context_call<'tcx>(bb_data: &mut BasicBlockData<'tcx>) -> Local {
@@ -1522,24 +1520,12 @@ impl<'tcx> crate::MirPass<'tcx> for StateTransform {
         // (finally in open_drop_for_tuple) before async drop expansion.
         // Async drops, produced by this drop elaboration, will be expanded,
         // and corresponding futures kept in layout.
-        let has_async_drops = matches!(
-            coroutine_kind,
-            CoroutineKind::Desugared(CoroutineDesugaring::Async | CoroutineDesugaring::AsyncGen, _)
-        ) && has_expandable_async_drops(tcx, body, coroutine_ty);
+        let coroutine_is_async = coroutine_kind.is_async_desugaring();
+        let has_async_drops = has_async_drops(body);
 
         // Replace all occurrences of `ResumeTy` with `&mut Context<'_>` within async bodies.
-        if matches!(
-            coroutine_kind,
-            CoroutineKind::Desugared(CoroutineDesugaring::Async | CoroutineDesugaring::AsyncGen, _)
-        ) {
-            let context_mut_ref = transform_async_context(tcx, body);
-            expand_async_drops(tcx, body, context_mut_ref, coroutine_kind, coroutine_ty);
-
-            if let Some(dumper) = MirDumper::new(tcx, "coroutine_async_drop_expand", body) {
-                dumper.dump_mir(body);
-            }
-        } else {
-            cleanup_async_drops(body);
+        if coroutine_is_async {
+            transform_async_context(tcx, body);
         }
 
         let always_live_locals = always_storage_live_locals(body);