Optimize async drop glue for some old types

Author: zetanumbers

compiler/rustc_hir/src/lang_items.rs

@@ -176,6 +176,7 @@ language_item_table! {
     AsyncDropSlice,          sym::async_drop_slice,    async_drop_slice_fn,        Target::Fn,             GenericRequirement::Exact(1);
     AsyncDropChain,          sym::async_drop_chain,    async_drop_chain_fn,        Target::Fn,             GenericRequirement::Exact(2);
     AsyncDropNoop,           sym::async_drop_noop,     async_drop_noop_fn,         Target::Fn,             GenericRequirement::Exact(0);
+    AsyncDropDeferredDropInPlace, sym::async_drop_deferred_drop_in_place, async_drop_deferred_drop_in_place_fn, Target::Fn, GenericRequirement::Exact(1);
     AsyncDropFuse,           sym::async_drop_fuse,     async_drop_fuse_fn,         Target::Fn,             GenericRequirement::Exact(1);
     AsyncDropDefer,          sym::async_drop_defer,    async_drop_defer_fn,        Target::Fn,             GenericRequirement::Exact(1);
     AsyncDropEither,         sym::async_drop_either,   async_drop_either_fn,       Target::Fn,             GenericRequirement::Exact(3);

compiler/rustc_middle/src/ty/adt.rs

@@ -579,8 +579,8 @@ impl<'tcx> AdtDef<'tcx> {
         tcx.adt_destructor(self.did())
     }
 
-    // FIXME(zetanumbers): consider supporting this method in same places where
-    // `destructor` is referenced
+    // FIXME: consider combining this method with `AdtDef::destructor` and removing
+    // this version
     pub fn async_destructor(self, tcx: TyCtxt<'tcx>) -> Option<AsyncDestructor> {
         tcx.adt_async_destructor(self.did())
     }

compiler/rustc_middle/src/ty/sty.rs

@@ -25,7 +25,7 @@ use std::assert_matches::debug_assert_matches;
 use std::borrow::Cow;
 use std::iter;
 use std::ops::{ControlFlow, Range};
-use ty::util::IntTypeExt;
+use ty::util::{AsyncDropGlueMorphology, IntTypeExt};
 
 use rustc_type_ir::TyKind::*;
 use rustc_type_ir::{self as ir, BoundVar, CollectAndApply, DynKind};
@@ -1934,11 +1934,22 @@ impl<'tcx> Ty<'tcx> {
     }
 
     /// Returns the type of the async destructor of this type.
-    pub fn async_destructor_ty(self, tcx: TyCtxt<'tcx>, param_env: ParamEnv<'tcx>) -> Ty<'tcx> {
-        if self.is_async_destructor_noop(tcx, param_env) || matches!(self.kind(), ty::Error(_)) {
-            return Ty::async_destructor_combinator(tcx, LangItem::AsyncDropNoop)
-                .instantiate_identity();
+    pub fn async_destructor_ty(self, tcx: TyCtxt<'tcx>) -> Ty<'tcx> {
+        match self.async_drop_glue_morphology(tcx) {
+            AsyncDropGlueMorphology::Noop => {
+                return Ty::async_destructor_combinator(tcx, LangItem::AsyncDropNoop)
+                    .instantiate_identity();
+            }
+            AsyncDropGlueMorphology::DeferredDropInPlace => {
+                let drop_in_place =
+                    Ty::async_destructor_combinator(tcx, LangItem::AsyncDropDeferredDropInPlace)
+                        .instantiate(tcx, &[self.into()]);
+                return Ty::async_destructor_combinator(tcx, LangItem::AsyncDropFuse)
+                    .instantiate(tcx, &[drop_in_place.into()]);
+            }
+            AsyncDropGlueMorphology::Custom => (),
         }
+
         match *self.kind() {
             ty::Param(_) | ty::Alias(..) | ty::Infer(ty::TyVar(_)) => {
                 let assoc_items = tcx
@@ -1957,19 +1968,13 @@ impl<'tcx> Ty<'tcx> {
                 .adt_async_destructor_ty(
                     tcx,
                     adt_def.variants().iter().map(|v| v.fields.iter().map(|f| f.ty(tcx, args))),
-                    param_env,
                 ),
-            ty::Tuple(tys) => self.adt_async_destructor_ty(tcx, iter::once(tys), param_env),
-            ty::Closure(_, args) => self.adt_async_destructor_ty(
-                tcx,
-                iter::once(args.as_closure().upvar_tys()),
-                param_env,
-            ),
-            ty::CoroutineClosure(_, args) => self.adt_async_destructor_ty(
-                tcx,
-                iter::once(args.as_coroutine_closure().upvar_tys()),
-                param_env,
-            ),
+            ty::Tuple(tys) => self.adt_async_destructor_ty(tcx, iter::once(tys)),
+            ty::Closure(_, args) => {
+                self.adt_async_destructor_ty(tcx, iter::once(args.as_closure().upvar_tys()))
+            }
+            ty::CoroutineClosure(_, args) => self
+                .adt_async_destructor_ty(tcx, iter::once(args.as_coroutine_closure().upvar_tys())),
 
             ty::Adt(adt_def, _) => {
                 assert!(adt_def.is_union());
@@ -1991,17 +1996,12 @@ impl<'tcx> Ty<'tcx> {
         }
     }
 
-    fn adt_async_destructor_ty<I>(
-        self,
-        tcx: TyCtxt<'tcx>,
-        variants: I,
-        param_env: ParamEnv<'tcx>,
-    ) -> Ty<'tcx>
+    fn adt_async_destructor_ty<I>(self, tcx: TyCtxt<'tcx>, variants: I) -> Ty<'tcx>
     where
         I: Iterator + ExactSizeIterator,
         I::Item: IntoIterator<Item = Ty<'tcx>>,
     {
-        debug_assert!(!self.is_async_destructor_noop(tcx, param_env));
+        debug_assert_eq!(self.async_drop_glue_morphology(tcx), AsyncDropGlueMorphology::Custom);
 
         let defer = Ty::async_destructor_combinator(tcx, LangItem::AsyncDropDefer);
         let chain = Ty::async_destructor_combinator(tcx, LangItem::AsyncDropChain);

compiler/rustc_middle/src/ty/util.rs

@@ -421,6 +421,25 @@ impl<'tcx> TyCtxt<'tcx> {
         Some(ty::AsyncDestructor { future, ctor })
     }
 
+    /// Returns async drop glue morphology for a definition. To get async drop
+    /// glue morphology for a type see [`Ty::async_drop_glue_morphology`].
+    //
+    // FIXME: consider making this a query
+    pub fn async_drop_glue_morphology(self, did: DefId) -> AsyncDropGlueMorphology {
+        let ty: Ty<'tcx> = self.type_of(did).instantiate_identity();
+
+        // Async drop glue morphology is an internal detail, so reveal_all probably
+        // should be fine
+        let param_env = ty::ParamEnv::reveal_all();
+        if ty.needs_async_drop(self, param_env) {
+            AsyncDropGlueMorphology::Custom
+        } else if ty.needs_drop(self, param_env) {
+            AsyncDropGlueMorphology::DeferredDropInPlace
+        } else {
+            AsyncDropGlueMorphology::Noop
+        }
+    }
+
     /// Returns the set of types that are required to be alive in
     /// order to run the destructor of `def` (see RFCs 769 and
     /// 1238).
@@ -1177,6 +1196,18 @@ impl<'tcx> TypeFolder<TyCtxt<'tcx>> for WeakAliasTypeExpander<'tcx> {
     }
 }
 
+/// Indicates the form of `AsyncDestruct::Destructor`. Used to simplify async
+/// drop glue for types not using async drop.
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub enum AsyncDropGlueMorphology {
+    /// Async destructor simply does nothing
+    Noop,
+    /// Async destructor simply runs `drop_in_place`
+    DeferredDropInPlace,
+    /// Async destructor has custom logic
+    Custom,
+}
+
 impl<'tcx> Ty<'tcx> {
     /// Returns the `Size` for primitive types (bool, uint, int, char, float).
     pub fn primitive_size(self, tcx: TyCtxt<'tcx>) -> Size {
@@ -1342,27 +1373,16 @@ impl<'tcx> Ty<'tcx> {
         }
     }
 
-    /// Checks whether values of this type `T` implement has noop async destructor.
+    /// Get morphology of the async drop glue, needed for types which do not
+    /// use async drop. To get async drop glue morphology for a definition see
+    /// [`TyCtxt::async_drop_glue_morphology`]. Used for `AsyncDestruct::Destructor`
+    /// type construction.
     //
-    // FIXME: implement optimization to make ADTs, which do not need drop,
-    // to skip fields or to have noop async destructor, use `needs_(async_)drop`
-    pub fn is_async_destructor_noop(
-        self,
-        tcx: TyCtxt<'tcx>,
-        param_env: ty::ParamEnv<'tcx>,
-    ) -> bool {
-        // TODO: check on the most generic version of your type
-        self.is_async_destructor_trivially_noop()
-            || self.needs_async_drop(tcx, param_env)
-            || self.needs_drop(tcx, param_env)
-    }
-
-    /// Fast path helper for testing if a type has noop async destructor.
-    ///
-    /// Returning `true` means the type is known to have noop async destructor
-    /// implementation. Returning `true` means nothing -- could be
-    /// `Drop`, might not be.
-    fn is_async_destructor_trivially_noop(self) -> bool {
+    // FIXME: implement optimization to not instantiate a certain morphology of
+    // async drop glue too soon to allow per type optimizations, see array case
+    // for more info. Perhaps then remove this method and use `needs_(async_)drop`
+    // instead.
+    pub fn async_drop_glue_morphology(self, tcx: TyCtxt<'tcx>) -> AsyncDropGlueMorphology {
         match self.kind() {
             ty::Int(_)
             | ty::Uint(_)
@@ -1374,37 +1394,43 @@ impl<'tcx> Ty<'tcx> {
             | ty::Ref(..)
             | ty::RawPtr(..)
             | ty::FnDef(..)
-            | ty::FnPtr(_) => true,
-            ty::Tuple(tys) => tys.is_empty(),
-            ty::Adt(adt_def, _) => adt_def.is_manually_drop(),
-            ty::Bool => todo!(),
-            ty::Char => todo!(),
-            ty::Int(_) => todo!(),
-            ty::Uint(_) => todo!(),
-            ty::Float(_) => todo!(),
-            ty::Adt(_, _) => todo!(),
-            ty::Foreign(_) => todo!(),
-            ty::Str => todo!(),
-            ty::Array(_, _) => todo!(),
-            ty::Pat(_, _) => todo!(),
-            ty::Slice(_) => todo!(),
-            ty::RawPtr(_, _) => todo!(),
-            ty::Ref(_, _, _) => todo!(),
-            ty::FnDef(_, _) => todo!(),
-            ty::FnPtr(_) => todo!(),
-            ty::Dynamic(_, _, _) => todo!(),
-            ty::Closure(_, _) => todo!(),
-            ty::CoroutineClosure(_, _) => todo!(),
-            ty::Coroutine(_, _) => todo!(),
-            ty::CoroutineWitness(_, _) => todo!(),
-            ty::Never => todo!(),
-            ty::Tuple(_) => todo!(),
-            ty::Alias(_, _) => todo!(),
-            ty::Param(_) => todo!(),
-            ty::Bound(_, _) => todo!(),
-            ty::Placeholder(_) => todo!(),
-            ty::Infer(_) => todo!(),
-            ty::Error(_) => todo!(),
+            | ty::FnPtr(_)
+            | ty::Infer(ty::FreshIntTy(_))
+            | ty::Infer(ty::FreshFloatTy(_)) => AsyncDropGlueMorphology::Noop,
+
+            ty::Tuple(tys) if tys.is_empty() => AsyncDropGlueMorphology::Noop,
+            ty::Adt(adt_def, _) if adt_def.is_manually_drop() => AsyncDropGlueMorphology::Noop,
+
+            // Foreign types can never have destructors.
+            ty::Foreign(_) => AsyncDropGlueMorphology::Noop,
+
+            // FIXME: implement dynamic types async drops
+            ty::Error(_) | ty::Dynamic(..) => AsyncDropGlueMorphology::DeferredDropInPlace,
+
+            ty::Tuple(_) | ty::Array(_, _) | ty::Slice(_) => {
+                // Assume worst-case scenario, because we can instantiate async
+                // destructors in different orders:
+                //
+                // 1. Instantiate [T; N] with T = String and N = 0
+                // 2. Instantiate <[String; 0] as AsyncDestruct>::Destructor
+                //
+                // And viceversa, thus we cannot rely on String not using async
+                // drop or array having zero (0) elements
+                AsyncDropGlueMorphology::Custom
+            }
+            ty::Pat(ty, _) => ty.async_drop_glue_morphology(tcx),
+
+            ty::Adt(adt_def, _) => tcx.async_drop_glue_morphology(adt_def.did()),
+
+            ty::Closure(did, _)
+            | ty::CoroutineClosure(did, _)
+            | ty::Coroutine(did, _)
+            | ty::CoroutineWitness(did, _) => tcx.async_drop_glue_morphology(*did),
+
+            ty::Alias(..) | ty::Param(_) | ty::Bound(..) | ty::Placeholder(..) | ty::Infer(_) => {
+                // No specifics, but would usually mean forwarding async drop glue
+                AsyncDropGlueMorphology::Custom
+            }
         }
     }
 
@@ -1451,7 +1477,11 @@ impl<'tcx> Ty<'tcx> {
     /// (Note that this implies that if `ty` has an async destructor attached,
     /// then `needs_async_drop` will definitely return `true` for `ty`.)
     ///
-    /// Note that this method is used to check eligible types in unions.
+    /// When constructing `AsyncDestruct::Destructor` type, use
+    /// [`Ty::async_drop_glue_morphology`] instead.
+    //
+    // FIXME(zetanumbers): Note that this method is used to check eligible types
+    // in unions.
     #[inline]
     pub fn needs_async_drop(self, tcx: TyCtxt<'tcx>, param_env: ty::ParamEnv<'tcx>) -> bool {
         // Avoid querying in simple cases.
@@ -1647,10 +1677,13 @@ impl<'tcx> ExplicitSelf<'tcx> {
     }
 }
 
-// FIXME(zetanumbers): make specifying asyncness explicit
 /// Returns a list of types such that the given type needs drop if and only if
 /// *any* of the returned types need drop. Returns `Err(AlwaysRequiresDrop)` if
 /// this type always needs drop.
+//
+// FIXME(zetanumbers): consider replacing this with only
+// `needs_drop_components_with_async`
+#[inline]
 pub fn needs_drop_components<'tcx>(
     tcx: TyCtxt<'tcx>,
     ty: Ty<'tcx>,

compiler/rustc_mir_transform/src/shim/async_destructor_ctor.rs

@@ -12,7 +12,7 @@ use rustc_middle::mir::{
     Terminator, TerminatorKind, UnwindAction, UnwindTerminateReason, RETURN_PLACE,
 };
 use rustc_middle::ty::adjustment::PointerCoercion;
-use rustc_middle::ty::util::Discr;
+use rustc_middle::ty::util::{AsyncDropGlueMorphology, Discr};
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_middle::{bug, span_bug};
 use rustc_span::source_map::respan;
@@ -116,15 +116,25 @@ impl<'tcx> AsyncDestructorCtorShimBuilder<'tcx> {
     }
 
     fn build(self) -> Body<'tcx> {
-        let (tcx, def_id, Some(self_ty)) = (self.tcx, self.def_id, self.self_ty) else {
+        let (tcx, Some(self_ty)) = (self.tcx, self.self_ty) else {
             return self.build_zst_output();
         };
+        match self_ty.async_drop_glue_morphology(tcx) {
+            AsyncDropGlueMorphology::Noop => span_bug!(
+                self.span,
+                "async drop glue shim generator encountered type with noop async drop glue morphology"
+            ),
+            AsyncDropGlueMorphology::DeferredDropInPlace => {
+                return self.build_deferred_drop_in_place();
+            }
+            AsyncDropGlueMorphology::Custom => (),
+        }
 
         let surface_drop_kind = || {
-            let param_env = tcx.param_env_reveal_all_normalized(def_id);
-            if self_ty.has_surface_async_drop(tcx, param_env) {
+            let adt_def = self_ty.ty_adt_def()?;
+            if adt_def.async_destructor(tcx).is_some() {
                 Some(SurfaceDropKind::Async)
-            } else if self_ty.has_surface_drop(tcx, param_env) {
+            } else if adt_def.destructor(tcx).is_some() {
                 Some(SurfaceDropKind::Sync)
             } else {
                 None
@@ -267,6 +277,13 @@ impl<'tcx> AsyncDestructorCtorShimBuilder<'tcx> {
         self.return_()
     }
 
+    fn build_deferred_drop_in_place(mut self) -> Body<'tcx> {
+        self.put_self();
+        let deferred = self.combine_deferred_drop_in_place();
+        self.combine_fuse(deferred);
+        self.return_()
+    }
+
     fn build_fused_async_surface(mut self) -> Body<'tcx> {
         self.put_self();
         let surface = self.combine_async_surface();
@@ -441,6 +458,14 @@ impl<'tcx> AsyncDestructorCtorShimBuilder<'tcx> {
         )
     }
 
+    fn combine_deferred_drop_in_place(&mut self) -> Ty<'tcx> {
+        self.apply_combinator(
+            1,
+            LangItem::AsyncDropDeferredDropInPlace,
+            &[self.self_ty.unwrap().into()],
+        )
+    }
+
     fn combine_fuse(&mut self, inner_future_ty: Ty<'tcx>) -> Ty<'tcx> {
         self.apply_combinator(1, LangItem::AsyncDropFuse, &[inner_future_ty.into()])
     }
@@ -481,7 +506,7 @@ impl<'tcx> AsyncDestructorCtorShimBuilder<'tcx> {
         if let Some(ty) = self.self_ty {
             debug_assert_eq!(
                 output.ty(&self.locals, self.tcx),
-                ty.async_destructor_ty(self.tcx, self.param_env),
+                ty.async_destructor_ty(self.tcx),
                 "output async destructor types did not match for type: {ty:?}",
             );
         }

compiler/rustc_span/src/symbol.rs

@@ -430,6 +430,7 @@ symbols! {
         async_drop,
         async_drop_chain,
         async_drop_defer,
+        async_drop_deferred_drop_in_place,
         async_drop_either,
         async_drop_fuse,
         async_drop_in_place,

compiler/rustc_trait_selection/src/solve/normalizes_to/mod.rs

@@ -831,7 +831,7 @@ impl<'tcx> assembly::GoalKind<'tcx> for NormalizesTo<'tcx> {
             | ty::Str
             | ty::Slice(_)
             | ty::Tuple(_)
-            | ty::Error(_) => self_ty.async_destructor_ty(ecx.tcx(), goal.param_env),
+            | ty::Error(_) => self_ty.async_destructor_ty(ecx.tcx()),
 
             // We do not call `Ty::async_destructor_ty` on alias, param, or placeholder
             // types, which return `<self_ty as AsyncDestruct>::AsyncDestructor`

compiler/rustc_trait_selection/src/traits/project.rs

@@ -1560,7 +1560,7 @@ fn confirm_builtin_candidate<'cx, 'tcx>(
         let destructor_def_id = tcx.associated_item_def_ids(trait_def_id)[0];
         assert_eq!(destructor_def_id, item_def_id);
 
-        (self_ty.async_destructor_ty(tcx, obligation.param_env).into(), Vec::new())
+        (self_ty.async_destructor_ty(tcx).into(), Vec::new())
     } else if lang_items.pointee_trait() == Some(trait_def_id) {
         let metadata_def_id = tcx.require_lang_item(LangItem::Metadata, None);
         assert_eq!(metadata_def_id, item_def_id);