Async drop support for dyn traits

Author: azhogin

compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs

@@ -1471,6 +1471,9 @@ fn build_vtable_type_di_node<'ll, 'tcx>(
                         ty::VtblEntry::MetadataDropInPlace => {
                             ("drop_in_place".to_string(), void_pointer_type_di_node)
                         }
+                        ty::VtblEntry::MetadataAsyncDropInPlace => {
+                            ("async_drop_in_place".to_string(), void_pointer_type_di_node)
+                        }
                         ty::VtblEntry::Method(_) => {
                             // Note: This code does not try to give a proper name to each method
                             //       because their might be multiple methods with the same name

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -610,7 +610,10 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 //                \-------/
                 //
                 let virtual_drop = Instance {
-                    def: ty::InstanceKind::Virtual(drop_fn.def_id(), 0), // idx 0: the drop function
+                    def: ty::InstanceKind::Virtual(
+                        drop_fn.def_id(),
+                        ty::COMMON_VTABLE_ENTRIES_DROPINPLACE,
+                    ),
                     args: drop_fn.args,
                 };
                 debug!("ty = {:?}", ty);
@@ -652,7 +655,10 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 //
                 // SO THEN WE CAN USE THE ABOVE CODE.
                 let virtual_drop = Instance {
-                    def: ty::InstanceKind::Virtual(drop_fn.def_id(), 0), // idx 0: the drop function
+                    def: ty::InstanceKind::Virtual(
+                        drop_fn.def_id(),
+                        ty::COMMON_VTABLE_ENTRIES_DROPINPLACE,
+                    ),
                     args: drop_fn.args,
                 };
                 debug!("ty = {:?}", ty);
@@ -707,6 +713,108 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         )
     }
 
+    #[tracing::instrument(level = "trace", skip(self, helper, bx))]
+    fn codegen_async_drop_dyn(
+        &mut self,
+        helper: TerminatorCodegenHelper<'tcx>,
+        bx: &mut Bx,
+        source_info: &mir::SourceInfo,
+        dropee: mir::Place<'tcx>,
+        destination: mir::Place<'tcx>,
+        target: mir::BasicBlock,
+        unwind: mir::UnwindAction,
+        mergeable_succ: bool,
+    ) -> MergingSucc {
+        let ty = dropee.ty(self.mir, bx.tcx()).ty;
+        let ty = self.monomorphize(ty);
+        let drop_fn = Instance::resolve_async_drop_in_place_dyn(bx.tcx(), ty).unwrap();
+        let place = self.codegen_place(bx, dropee.as_ref());
+
+        let (args1, args2);
+        let mut args = if let Some(llextra) = place.val.llextra {
+            args2 = [place.val.llval, llextra];
+            &args2[..]
+        } else {
+            args1 = [place.val.llval];
+            &args1[..]
+        };
+
+        let (drop_fn, fn_abi, drop_instance) = match ty.kind() {
+            ty::Dynamic(_, _, ty::Dyn) => {
+                let virtual_drop = Instance {
+                    def: ty::InstanceKind::Virtual(
+                        drop_fn.def_id(),
+                        ty::COMMON_VTABLE_ENTRIES_ASYNCDROPINPLACE,
+                    ),
+                    args: drop_fn.args,
+                };
+                let fn_abi = bx.fn_abi_of_instance(virtual_drop, ty::List::empty());
+                let vtable = args[1];
+                // Truncate vtable off of args list
+                args = &args[..1];
+                (
+                    meth::VirtualIndex::from_index(ty::COMMON_VTABLE_ENTRIES_ASYNCDROPINPLACE)
+                        .get_optional_fn(bx, vtable, ty, fn_abi),
+                    fn_abi,
+                    virtual_drop,
+                )
+            }
+            ty::Dynamic(_, _, ty::DynStar) => {
+                let virtual_drop = Instance {
+                    def: ty::InstanceKind::Virtual(
+                        drop_fn.def_id(),
+                        ty::COMMON_VTABLE_ENTRIES_ASYNCDROPINPLACE,
+                    ),
+                    args: drop_fn.args,
+                };
+                let fn_abi = bx.fn_abi_of_instance(virtual_drop, ty::List::empty());
+                let meta_ptr = place.project_field(bx, 1);
+                let meta = bx.load_operand(meta_ptr);
+                // Truncate vtable off of args list
+                assert!(args.len() == 1);
+                args = &args[..1];
+                debug!("args' = {:?}", args);
+                (
+                    meth::VirtualIndex::from_index(ty::COMMON_VTABLE_ENTRIES_ASYNCDROPINPLACE)
+                        .get_optional_fn(bx, meta.immediate(), ty, fn_abi),
+                    fn_abi,
+                    virtual_drop,
+                )
+            }
+            _ => bug!("Non-virtual call for async drop terminator (ty is not dyn or dyn*)"),
+        };
+        // We generate a null check for the drop_fn. This saves a bunch of relocations being
+        // generated for no-op drops.
+        // FIXME: do we need it for dyn async drop?
+        {
+            let is_not_null = bx.append_sibling_block("is_not_null");
+            let llty = bx.fn_ptr_backend_type(fn_abi);
+            let null = bx.const_null(llty);
+            let non_null =
+                bx.icmp(base::bin_op_to_icmp_predicate(mir::BinOp::Ne, false), drop_fn, null);
+            bx.cond_br(non_null, is_not_null, helper.llbb_with_cleanup(self, target));
+            bx.switch_to_block(is_not_null);
+            self.set_debug_loc(bx, *source_info);
+        }
+        assert!(!fn_abi.ret.is_indirect());
+        let mut llargs = Vec::new();
+        let return_dest = self.make_return_dest(bx, destination, &fn_abi.ret, &mut llargs);
+        assert!(llargs.is_empty());
+
+        helper.do_call(
+            self,
+            bx,
+            fn_abi,
+            drop_fn,
+            args,
+            Some((return_dest, target)),
+            unwind,
+            &[],
+            Some(drop_instance),
+            false,
+        )
+    }
+
     fn codegen_assert_terminator(
         &mut self,
         helper: TerminatorCodegenHelper<'tcx>,
@@ -907,6 +1015,21 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
 
         let (instance, mut llfn) = match *callee.layout.ty.kind() {
             ty::FnDef(def_id, generic_args) => {
+                if bx.tcx().is_lang_item(def_id, LangItem::AsyncDropInPlaceDyn) {
+                    let mir::Operand::Move(dropee) = args[0].node else {
+                        bug!();
+                    };
+                    return self.codegen_async_drop_dyn(
+                        helper,
+                        bx,
+                        &terminator.source_info,
+                        dropee,
+                        destination,
+                        target.unwrap(),
+                        unwind,
+                        mergeable_succ,
+                    );
+                }
                 let instance = ty::Instance::expect_resolve(
                     bx.tcx(),
                     bx.typing_env(),

compiler/rustc_const_eval/src/interpret/call.rs

@@ -825,6 +825,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         // only supports calling `VtblEntry::Method`; it would choke on a `MetadataDropInPlace`. So
         // instead we do the virtual call stuff ourselves. It's easier here than in `eval_fn_call`
         // since we can just get a place of the underlying type and use `mplace_to_ref`.
+        // FIXME: Support AsyncDrop (async_drop_in_place) here
         let place = match place.layout.ty.kind() {
             ty::Dynamic(data, _, ty::Dyn) => {
                 // Dropping a trait object. Need to find actual drop fn.

compiler/rustc_hir/src/lang_items.rs

@@ -191,6 +191,8 @@ language_item_table! {
     Destruct,                sym::destruct,            destruct_trait,             Target::Trait,          GenericRequirement::None;
     AsyncDrop,               sym::async_drop,          async_drop_trait,           Target::Trait,          GenericRequirement::None;
     AsyncDropInPlace,        sym::async_drop_in_place, async_drop_in_place_fn,     Target::Fn,             GenericRequirement::Exact(1);
+    AsyncDropInPlaceDyn,     sym::async_drop_in_place_dyn, async_drop_in_place_dyn_fn, Target::Fn,         GenericRequirement::Exact(1);
+    AsyncDropInPlaceSelf,    sym::async_drop_in_place_self, async_drop_in_place_self_fn, Target::Fn,       GenericRequirement::None;
 
     CoerceUnsized,           sym::coerce_unsized,      coerce_unsized_trait,       Target::Trait,          GenericRequirement::Minimum(1);
     DispatchFromDyn,         sym::dispatch_from_dyn,   dispatch_from_dyn_trait,    Target::Trait,          GenericRequirement::Minimum(1);

compiler/rustc_middle/src/ty/instance.rs

@@ -809,6 +809,30 @@ impl<'tcx> Instance<'tcx> {
         )
     }
 
+    // async_drop_in_place, with return value converted into Pin<Box<Future>>, for usage in vtable.
+    // May returns None for core lib compilation (before lang item definition in alloc lib).
+    pub fn resolve_async_drop_in_place_dyn(
+        tcx: TyCtxt<'tcx>,
+        ty: Ty<'tcx>,
+    ) -> Option<ty::Instance<'tcx>> {
+        // `async_drop_in_place<T>::{closure}` is a special case, because such coroutine is its async drop future itself.
+        // To drop this coroutine we need to continue poll it.
+        // So, it async drop constructor function in vtable returns its address (from argument), boxed and pinned.
+        let (item, args) = if ty.is_async_drop_in_place_coroutine(tcx) {
+            (LangItem::AsyncDropInPlaceSelf, tcx.mk_args(&[]))
+        } else {
+            (LangItem::AsyncDropInPlaceDyn, tcx.mk_args(&[ty.into()]))
+        };
+        let Some(def_id) = tcx.lang_items().get(item) else { return None };
+        Some(Instance::expect_resolve(
+            tcx,
+            ty::TypingEnv::fully_monomorphized(),
+            def_id,
+            args,
+            ty.ty_adt_def().and_then(|adt| tcx.hir_span_if_local(adt.did())).unwrap_or(DUMMY_SP),
+        ))
+    }
+
     pub fn resolve_async_drop_in_place_poll(
         tcx: TyCtxt<'tcx>,
         def_id: DefId,

compiler/rustc_middle/src/ty/vtable.rs

@@ -13,6 +13,8 @@ use crate::ty::{self, Instance, TraitRef, Ty, TyCtxt};
 pub enum VtblEntry<'tcx> {
     /// destructor of this type (used in vtable header)
     MetadataDropInPlace,
+    /// async destructor of this type, function returns Pin<Box<Future>>> (used in vtable header)
+    MetadataAsyncDropInPlace,
     /// layout size of this type (used in vtable header)
     MetadataSize,
     /// layout align of this type (used in vtable header)
@@ -31,6 +33,7 @@ impl<'tcx> fmt::Debug for VtblEntry<'tcx> {
         // so we implement this manually.
         match self {
             VtblEntry::MetadataDropInPlace => write!(f, "MetadataDropInPlace"),
+            VtblEntry::MetadataAsyncDropInPlace => write!(f, "MetadataAsyncDropInPlace"),
             VtblEntry::MetadataSize => write!(f, "MetadataSize"),
             VtblEntry::MetadataAlign => write!(f, "MetadataAlign"),
             VtblEntry::Vacant => write!(f, "Vacant"),
@@ -42,12 +45,17 @@ impl<'tcx> fmt::Debug for VtblEntry<'tcx> {
 
 // Needs to be associated with the `'tcx` lifetime
 impl<'tcx> TyCtxt<'tcx> {
-    pub const COMMON_VTABLE_ENTRIES: &'tcx [VtblEntry<'tcx>] =
-        &[VtblEntry::MetadataDropInPlace, VtblEntry::MetadataSize, VtblEntry::MetadataAlign];
+    pub const COMMON_VTABLE_ENTRIES: &'tcx [VtblEntry<'tcx>] = &[
+        VtblEntry::MetadataDropInPlace,
+        VtblEntry::MetadataAsyncDropInPlace,
+        VtblEntry::MetadataSize,
+        VtblEntry::MetadataAlign,
+    ];
 }
 
 pub const COMMON_VTABLE_ENTRIES_DROPINPLACE: usize = 0;
-pub const COMMON_VTABLE_ENTRIES_SIZE: usize = 1;
+pub const COMMON_VTABLE_ENTRIES_ASYNCDROPINPLACE: usize = 1;
+pub const COMMON_VTABLE_ENTRIES_SIZE: usize = 2;
 pub const COMMON_VTABLE_ENTRIES_ALIGN: usize = 2;
 
 // Note that we don't have access to a self type here, this has to be purely based on the trait (and
@@ -129,6 +137,18 @@ pub(super) fn vtable_allocation_provider<'tcx>(
                     Scalar::from_maybe_pointer(Pointer::null(), &tcx)
                 }
             }
+            VtblEntry::MetadataAsyncDropInPlace => {
+                if tcx.features().async_drop()
+                    && ty.needs_async_drop(tcx, ty::TypingEnv::fully_monomorphized())
+                    && let Some(instance) = ty::Instance::resolve_async_drop_in_place_dyn(tcx, ty)
+                {
+                    let fn_alloc_id = tcx.reserve_and_set_fn_alloc(instance, CTFE_ALLOC_SALT);
+                    let fn_ptr = Pointer::from(fn_alloc_id);
+                    Scalar::from_pointer(fn_ptr, &tcx)
+                } else {
+                    Scalar::from_maybe_pointer(Pointer::null(), &tcx)
+                }
+            }
             VtblEntry::MetadataSize => Scalar::from_uint(size, ptr_size),
             VtblEntry::MetadataAlign => Scalar::from_uint(align, ptr_size),
             VtblEntry::Vacant => continue,