Auto merge of rust-lang#124662 - zetanumbers:needs_async_drop, r=oli-obk

Implement `needs_async_drop` in rustc and optimize async drop glue

This PR expands on rust-lang#121801 and implements `Ty::needs_async_drop` which works almost exactly the same as `Ty::needs_drop`, which is needed for rust-lang#123948.

Also made compiler's async drop code to look more like compiler's regular drop code, which enabled me to write an optimization where types which do not use `AsyncDrop` can simply forward async drop glue to `drop_in_place`. This made size of the async block from the [async_drop test](https://github.com/zetanumbers/rust/blob/67980dd6fb11917d23d01a19c2cf4cfc3978aac8/tests/ui/async-await/async-drop.rs) to decrease by 12%.

Author: bors

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_hir_analysis/src/check/mod.rs

@@ -112,6 +112,7 @@ pub fn provide(providers: &mut Providers) {
     wfcheck::provide(providers);
     *providers = Providers {
         adt_destructor,
+        adt_async_destructor,
         region_scope_tree,
         collect_return_position_impl_trait_in_trait_tys,
         compare_impl_const: compare_impl_item::compare_impl_const_raw,
@@ -124,6 +125,10 @@ fn adt_destructor(tcx: TyCtxt<'_>, def_id: LocalDefId) -> Option<ty::Destructor>
     tcx.calculate_dtor(def_id.to_def_id(), dropck::check_drop_impl)
 }
 
+fn adt_async_destructor(tcx: TyCtxt<'_>, def_id: LocalDefId) -> Option<ty::AsyncDestructor> {
+    tcx.calculate_async_dtor(def_id.to_def_id(), dropck::check_drop_impl)
+}
+
 /// Given a `DefId` for an opaque type in return position, find its parent item's return
 /// expressions.
 fn get_owner_return_paths(

compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs

@@ -287,6 +287,10 @@ provide! { tcx, def_id, other, cdata,
         let _ = cdata;
         tcx.calculate_dtor(def_id, |_,_| Ok(()))
     }
+    adt_async_destructor => {
+        let _ = cdata;
+        tcx.calculate_async_dtor(def_id, |_,_| Ok(()))
+    }
     associated_item_def_ids => {
         tcx.arena.alloc_from_iter(cdata.get_associated_item_or_field_def_ids(def_id.index))
     }

compiler/rustc_middle/src/query/erase.rs

@@ -238,6 +238,7 @@ trivial! {
     Option<rustc_hir::CoroutineKind>,
     Option<rustc_hir::HirId>,
     Option<rustc_middle::middle::stability::DeprecationEntry>,
+    Option<rustc_middle::ty::AsyncDestructor>,
     Option<rustc_middle::ty::Destructor>,
     Option<rustc_middle::ty::ImplTraitInTraitData>,
     Option<rustc_middle::ty::ScalarInt>,
@@ -295,6 +296,7 @@ trivial! {
     rustc_middle::ty::AssocItem,
     rustc_middle::ty::AssocItemContainer,
     rustc_middle::ty::Asyncness,
+    rustc_middle::ty::AsyncDestructor,
     rustc_middle::ty::BoundVariableKind,
     rustc_middle::ty::DeducedParamAttrs,
     rustc_middle::ty::Destructor,

compiler/rustc_middle/src/query/mod.rs

@@ -703,6 +703,11 @@ rustc_queries! {
         cache_on_disk_if { key.is_local() }
         separate_provide_extern
     }
+    query adt_async_destructor(key: DefId) -> Option<ty::AsyncDestructor> {
+        desc { |tcx| "computing `AsyncDrop` impl for `{}`", tcx.def_path_str(key) }
+        cache_on_disk_if { key.is_local() }
+        separate_provide_extern
+    }
 
     query adt_sized_constraint(key: DefId) -> Option<ty::EarlyBinder<'tcx, Ty<'tcx>>> {
         desc { |tcx| "computing the `Sized` constraint for `{}`", tcx.def_path_str(key) }
@@ -1343,18 +1348,14 @@ rustc_queries! {
     query is_unpin_raw(env: ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> bool {
         desc { "computing whether `{}` is `Unpin`", env.value }
     }
-    /// Query backing `Ty::has_surface_async_drop`.
-    query has_surface_async_drop_raw(env: ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> bool {
-        desc { "computing whether `{}` has `AsyncDrop` implementation", env.value }
-    }
-    /// Query backing `Ty::has_surface_drop`.
-    query has_surface_drop_raw(env: ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> bool {
-        desc { "computing whether `{}` has `Drop` implementation", env.value }
-    }
     /// Query backing `Ty::needs_drop`.
     query needs_drop_raw(env: ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> bool {
         desc { "computing whether `{}` needs drop", env.value }
     }
+    /// Query backing `Ty::needs_async_drop`.
+    query needs_async_drop_raw(env: ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> bool {
+        desc { "computing whether `{}` needs async drop", env.value }
+    }
     /// Query backing `Ty::has_significant_drop_raw`.
     query has_significant_drop_raw(env: ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> bool {
         desc { "computing whether `{}` has a significant drop", env.value }

compiler/rustc_middle/src/ty/adt.rs

@@ -24,7 +24,9 @@ use std::hash::{Hash, Hasher};
 use std::ops::Range;
 use std::str;
 
-use super::{Destructor, FieldDef, GenericPredicates, Ty, TyCtxt, VariantDef, VariantDiscr};
+use super::{
+    AsyncDestructor, Destructor, FieldDef, GenericPredicates, Ty, TyCtxt, VariantDef, VariantDiscr,
+};
 
 #[derive(Clone, Copy, PartialEq, Eq, Hash, HashStable, TyEncodable, TyDecodable)]
 pub struct AdtFlags(u16);
@@ -577,6 +579,12 @@ impl<'tcx> AdtDef<'tcx> {
         tcx.adt_destructor(self.did())
     }
 
+    // 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())
+    }
+
     /// Returns a type such that `Self: Sized` if and only if that type is `Sized`,
     /// or `None` if the type is always sized.
     pub fn sized_constraint(self, tcx: TyCtxt<'tcx>) -> Option<ty::EarlyBinder<'tcx, Ty<'tcx>>> {

compiler/rustc_middle/src/ty/mod.rs

@@ -1172,6 +1172,15 @@ pub struct Destructor {
     pub constness: hir::Constness,
 }
 
+// FIXME: consider combining this definition with regular `Destructor`
+#[derive(Copy, Clone, Debug, HashStable, Encodable, Decodable)]
+pub struct AsyncDestructor {
+    /// The `DefId` of the async destructor future constructor
+    pub ctor: DefId,
+    /// The `DefId` of the async destructor future type
+    pub future: DefId,
+}
+
 #[derive(Clone, Copy, PartialEq, Eq, HashStable, TyEncodable, TyDecodable)]
 pub struct VariantFlags(u8);
 bitflags::bitflags! {

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};
@@ -1951,11 +1951,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
@@ -1974,24 +1985,18 @@ 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());
 
-                let surface_drop = self.surface_async_dropper_ty(tcx, param_env).unwrap();
+                let surface_drop = self.surface_async_dropper_ty(tcx).unwrap();
 
                 Ty::async_destructor_combinator(tcx, LangItem::AsyncDropFuse)
                     .instantiate(tcx, &[surface_drop.into()])
@@ -2008,17 +2013,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);
@@ -2041,7 +2041,7 @@ impl<'tcx> Ty<'tcx> {
             })
             .unwrap();
 
-        let dtor = if let Some(dropper_ty) = self.surface_async_dropper_ty(tcx, param_env) {
+        let dtor = if let Some(dropper_ty) = self.surface_async_dropper_ty(tcx) {
             Ty::async_destructor_combinator(tcx, LangItem::AsyncDropChain)
                 .instantiate(tcx, &[dropper_ty.into(), variants_dtor.into()])
         } else {
@@ -2052,21 +2052,13 @@ impl<'tcx> Ty<'tcx> {
             .instantiate(tcx, &[dtor.into()])
     }
 
-    fn surface_async_dropper_ty(
-        self,
-        tcx: TyCtxt<'tcx>,
-        param_env: ParamEnv<'tcx>,
-    ) -> Option<Ty<'tcx>> {
-        if self.has_surface_async_drop(tcx, param_env) {
-            Some(LangItem::SurfaceAsyncDropInPlace)
-        } else if self.has_surface_drop(tcx, param_env) {
-            Some(LangItem::AsyncDropSurfaceDropInPlace)
-        } else {
-            None
-        }
-        .map(|dropper| {
-            Ty::async_destructor_combinator(tcx, dropper).instantiate(tcx, &[self.into()])
-        })
+    fn surface_async_dropper_ty(self, tcx: TyCtxt<'tcx>) -> Option<Ty<'tcx>> {
+        let adt_def = self.ty_adt_def()?;
+        let dropper = adt_def
+            .async_destructor(tcx)
+            .map(|_| LangItem::SurfaceAsyncDropInPlace)
+            .or_else(|| adt_def.destructor(tcx).map(|_| LangItem::AsyncDropSurfaceDropInPlace))?;
+        Some(Ty::async_destructor_combinator(tcx, dropper).instantiate(tcx, &[self.into()]))
     }
 
     fn async_destructor_combinator(