Async drop codegen (WIP)

compiler/rustc_borrowck/src/lib.rs

@@ -675,7 +675,7 @@ impl<'cx, 'tcx, R> rustc_mir_dataflow::ResultsVisitor<'cx, 'tcx, R> for MirBorro
             TerminatorKind::SwitchInt { discr, targets: _ } => {
                 self.consume_operand(loc, (discr, span), flow_state);
             }
-            TerminatorKind::Drop { place, target: _, unwind: _, replace } => {
+            TerminatorKind::Drop { place, target: _, unwind: _, replace, drop: _, async_fut: _ } => {
                 debug!(
                     "visit_terminator_drop \
                      loc: {:?} term: {:?} place: {:?} span: {:?}",

compiler/rustc_borrowck/src/polonius/loan_invalidations.rs

@@ -99,7 +99,7 @@ impl<'cx, 'tcx> Visitor<'tcx> for LoanInvalidationsGenerator<'cx, 'tcx> {
             TerminatorKind::SwitchInt { discr, targets: _ } => {
                 self.consume_operand(location, discr);
             }
-            TerminatorKind::Drop { place: drop_place, target: _, unwind: _, replace } => {
+            TerminatorKind::Drop { place: drop_place, target: _, unwind: _, replace, drop: _, async_fut: _ } => {
                 let write_kind =
                     if *replace { WriteKind::Replace } else { WriteKind::StorageDeadOrDrop };
                 self.access_place(

compiler/rustc_borrowck/src/type_check/mod.rs

@@ -1684,8 +1684,14 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                 }
             }
             TerminatorKind::Unreachable => {}
-            TerminatorKind::Drop { target, unwind, .. }
-            | TerminatorKind::Assert { target, unwind, .. } => {
+            TerminatorKind::Drop { target, unwind, drop, .. } => {
+                self.assert_iscleanup(body, block_data, target, is_cleanup);
+                self.assert_iscleanup_unwind(body, block_data, unwind, is_cleanup);
+                if let Some(drop) = drop {
+                    self.assert_iscleanup(body, block_data, drop, is_cleanup);
+                }
+            }
+            TerminatorKind::Assert { target, unwind, .. } => {
                 self.assert_iscleanup(body, block_data, target, is_cleanup);
                 self.assert_iscleanup_unwind(body, block_data, unwind, is_cleanup);
             }

compiler/rustc_codegen_cranelift/src/base.rs

@@ -481,7 +481,7 @@ fn codegen_fn_body(fx: &mut FunctionCx<'_, '_, '_>, start_block: Block) {
             | TerminatorKind::CoroutineDrop => {
                 bug!("shouldn't exist at codegen {:?}", bb_data.terminator());
             }
-            TerminatorKind::Drop { place, target, unwind: _, replace: _ } => {
+            TerminatorKind::Drop { place, target, unwind: _, replace: _, drop: _, async_fut: _ } => {
                 let drop_place = codegen_place(fx, *place);
                 crate::abi::codegen_drop(fx, source_info, drop_place);
 

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -1232,7 +1232,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 MergingSucc::False
             }
 
-            mir::TerminatorKind::Drop { place, target, unwind, replace: _ } => {
+            mir::TerminatorKind::Drop { place, target, unwind, replace: _, drop: _, async_fut: _ } => {
                 self.codegen_drop_terminator(helper, bx, place, target, unwind, mergeable_succ())
             }
 

compiler/rustc_const_eval/src/const_eval/machine.rs

@@ -568,6 +568,7 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for CompileTimeInterpreter<'mir,
             RemainderByZero(op) => RemainderByZero(eval_to_int(op)?),
             ResumedAfterReturn(coroutine_kind) => ResumedAfterReturn(*coroutine_kind),
             ResumedAfterPanic(coroutine_kind) => ResumedAfterPanic(*coroutine_kind),
+            ResumedAfterDrop(coroutine_kind) => ResumedAfterDrop(*coroutine_kind),
             MisalignedPointerDereference { ref required, ref found } => {
                 MisalignedPointerDereference {
                     required: eval_to_int(required)?,

compiler/rustc_const_eval/src/interpret/terminator.rs

@@ -170,7 +170,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 }
             }
 
-            Drop { place, target, unwind, replace: _ } => {
+            Drop { place, target, unwind, replace: _, drop: _, async_fut: _ } => {
                 let frame = self.frame();
                 let ty = place.ty(&frame.body.local_decls, *self.tcx).ty;
                 let ty = self.subst_from_frame_and_normalize_erasing_regions(frame, ty)?;
@@ -542,6 +542,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             | ty::InstanceDef::ReifyShim(..)
             | ty::InstanceDef::ClosureOnceShim { .. }
             | ty::InstanceDef::FnPtrShim(..)
+            | ty::InstanceDef::FutureDropPollShim(..)
             | ty::InstanceDef::DropGlue(..)
             | ty::InstanceDef::CloneShim(..)
             | ty::InstanceDef::FnPtrAddrShim(..)

compiler/rustc_const_eval/src/transform/validate.rs

@@ -363,9 +363,12 @@ impl<'a, 'tcx> Visitor<'tcx> for CfgChecker<'a, 'tcx> {
                     );
                 }
             }
-            TerminatorKind::Drop { target, unwind, .. } => {
+            TerminatorKind::Drop { target, unwind, drop, .. } => {
                 self.check_edge(location, *target, EdgeKind::Normal);
                 self.check_unwind_edge(location, *unwind);
+                if let Some(drop) = drop {
+                    self.check_edge(location, *drop, EdgeKind::Normal);
+                }
             }
             TerminatorKind::Call { args, destination, target, unwind, .. } => {
                 if let Some(target) = target {

compiler/rustc_feature/src/unstable.rs

@@ -355,6 +355,8 @@ declare_features! (
     (unstable, associated_type_defaults, "1.2.0", Some(29661)),
     /// Allows `async || body` closures.
     (unstable, async_closure, "1.37.0", Some(62290)),
+    /// Allows implementing `AsyncDrop`.
+    (incomplete, async_drop, "CURRENT_RUSTC_VERSION", None),
     /// Allows `#[track_caller]` on async functions.
     (unstable, async_fn_track_caller, "1.73.0", Some(110011)),
     /// Allows `for await` loops.

compiler/rustc_hir/src/lang_items.rs

@@ -163,6 +163,9 @@ language_item_table! {
 
     Drop,                    sym::drop,                drop_trait,                 Target::Trait,          GenericRequirement::None;
     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);
+    FutureDropPoll,          sym::future_drop_poll,    future_drop_poll_fn,        Target::Fn,             GenericRequirement::Exact(1);
 
     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);
@@ -260,6 +263,7 @@ language_item_table! {
 
     ExchangeMalloc,          sym::exchange_malloc,     exchange_malloc_fn,         Target::Fn,             GenericRequirement::None;
     DropInPlace,             sym::drop_in_place,       drop_in_place_fn,           Target::Fn,             GenericRequirement::Minimum(1);
+    DropInPlaceFuture,       sym::drop_in_place_future,drop_in_place_future_fn,    Target::Fn,             GenericRequirement::Minimum(1);
     AllocLayout,             sym::alloc_layout,        alloc_layout,               Target::Struct,         GenericRequirement::None;
 
     Start,                   sym::start,               start_fn,                   Target::Fn,             GenericRequirement::Exact(1);

compiler/rustc_middle/messages.ftl

@@ -5,15 +5,21 @@ middle_assert_async_resume_after_panic = `async fn` resumed after panicking
 
 middle_assert_async_resume_after_return = `async fn` resumed after completion
 
+middle_assert_async_resume_after_drop = `async fn` resumed after async drop
+
 middle_assert_coroutine_resume_after_panic = coroutine resumed after panicking
 
 middle_assert_coroutine_resume_after_return = coroutine resumed after completion
 
+middle_assert_coroutine_resume_after_drop = coroutine resumed after async drop
+
 middle_assert_divide_by_zero =
     attempt to divide `{$val}` by zero
 
 middle_assert_gen_resume_after_panic = `gen` fn or block cannot be further iterated on after it panicked
 
+middle_assert_gen_resume_after_drop = `gen` fn or block cannot be further iterated on after it async dropped
+
 middle_assert_misaligned_ptr_deref =
     misaligned pointer dereference: address must be a multiple of {$required} but is {$found}
 

compiler/rustc_middle/src/mir/mod.rs

@@ -255,6 +255,9 @@ pub struct CoroutineInfo<'tcx> {
     /// Coroutine drop glue.
     pub coroutine_drop: Option<Body<'tcx>>,
 
+    /// Coroutine async drop glue.
+    pub coroutine_drop_async: Option<Body<'tcx>>,
+
     /// The layout of a coroutine. Produced by the state transformation.
     pub coroutine_layout: Option<CoroutineLayout<'tcx>>,
 
@@ -275,6 +278,7 @@ impl<'tcx> CoroutineInfo<'tcx> {
             yield_ty: Some(yield_ty),
             resume_ty: Some(resume_ty),
             coroutine_drop: None,
+            coroutine_drop_async: None,
             coroutine_layout: None,
         }
     }
@@ -578,6 +582,11 @@ impl<'tcx> Body<'tcx> {
         self.coroutine.as_ref().and_then(|coroutine| coroutine.coroutine_drop.as_ref())
     }
 
+    #[inline]
+    pub fn coroutine_drop_async(&self) -> Option<&Body<'tcx>> {
+        self.coroutine.as_ref().and_then(|coroutine| coroutine.coroutine_drop_async.as_ref())
+    }
+
     #[inline]
     pub fn coroutine_kind(&self) -> Option<CoroutineKind> {
         self.coroutine.as_ref().map(|coroutine| coroutine.coroutine_kind)

compiler/rustc_middle/src/mir/mono.rs

@@ -402,6 +402,7 @@ impl<'tcx> CodegenUnit<'tcx> {
                             | InstanceDef::FnPtrShim(..)
                             | InstanceDef::Virtual(..)
                             | InstanceDef::ClosureOnceShim { .. }
+                            | InstanceDef::FutureDropPollShim(..)
                             | InstanceDef::DropGlue(..)
                             | InstanceDef::CloneShim(..)
                             | InstanceDef::ThreadLocalShim(..)

compiler/rustc_middle/src/mir/pretty.rs

@@ -855,7 +855,9 @@ impl<'tcx> TerminatorKind<'tcx> {
             Call { target: None, unwind: _, .. } => vec![],
             Yield { drop: Some(_), .. } => vec!["resume".into(), "drop".into()],
             Yield { drop: None, .. } => vec!["resume".into()],
-            Drop { unwind: UnwindAction::Cleanup(_), .. } => vec!["return".into(), "unwind".into()],
+            Drop { unwind: UnwindAction::Cleanup(_), drop: Some(_), .. } => vec!["return".into(), "unwind".into(), "drop".into() ],
+            Drop { unwind: UnwindAction::Cleanup(_), drop: None, .. } => vec!["return".into(), "unwind".into()],
+            Drop { unwind: _, drop: Some(_), .. } => vec!["return".into(), "drop".into() ],
             Drop { unwind: _, .. } => vec!["return".into()],
             Assert { unwind: UnwindAction::Cleanup(_), .. } => {
                 vec!["success".into(), "unwind".into()]

compiler/rustc_middle/src/mir/syntax.rs

@@ -656,7 +656,16 @@ pub enum TerminatorKind<'tcx> {
     /// The `replace` flag indicates whether this terminator was created as part of an assignment.
     /// This should only be used for diagnostic purposes, and does not have any operational
     /// meaning.
-    Drop { place: Place<'tcx>, target: BasicBlock, unwind: UnwindAction, replace: bool },
+    Drop {
+        place: Place<'tcx>,
+        target: BasicBlock,
+        unwind: UnwindAction,
+        replace: bool,
+        /// Cleanup to be done if the coroutine is dropped at this suspend point (for async drop).
+        drop: Option<BasicBlock>,
+        /// Prepared async future local (for async drop)
+        async_fut: Option<Local>,
+    },
 
     /// Roughly speaking, evaluates the `func` operand and the arguments, and starts execution of
     /// the referred to function. The operand types must match the argument types of the function.
@@ -888,6 +897,7 @@ pub enum AssertKind<O> {
     RemainderByZero(O),
     ResumedAfterReturn(CoroutineKind),
     ResumedAfterPanic(CoroutineKind),
+    ResumedAfterDrop(CoroutineKind),
     MisalignedPointerDereference { required: O, found: O },
 }