Prereq5 for async drop - AsyncDropGlue & FutureDropPoll instances preparation

Author: azhogin

compiler/rustc_codegen_ssa/src/back/symbol_export.rs

@@ -364,7 +364,7 @@ fn exported_symbols_provider_local(
                     ));
                 }
                 MonoItem::Fn(Instance {
-                    def: InstanceKind::AsyncDropGlueCtorShim(def_id, Some(ty)),
+                    def: InstanceKind::AsyncDropGlueCtorShim(def_id, ty),
                     args,
                 }) => {
                     // A little sanity-check
@@ -381,6 +381,16 @@ fn exported_symbols_provider_local(
                         },
                     ));
                 }
+                MonoItem::Fn(Instance { def: InstanceKind::AsyncDropGlue(_, ty), args: _ }) => {
+                    symbols.push((
+                        ExportedSymbol::AsyncDropGlue(ty),
+                        SymbolExportInfo {
+                            level: SymbolExportLevel::Rust,
+                            kind: SymbolExportKind::Text,
+                            used: false,
+                        },
+                    ));
+                }
                 _ => {
                     // Any other symbols don't qualify for sharing
                 }
@@ -404,6 +414,7 @@ fn upstream_monomorphizations_provider(
 
     let drop_in_place_fn_def_id = tcx.lang_items().drop_in_place_fn();
     let async_drop_in_place_fn_def_id = tcx.lang_items().async_drop_in_place_fn();
+    let async_drop_in_place_poll_fn_def_id = tcx.lang_items().async_drop_in_place_poll_fn();
 
     for &cnum in cnums.iter() {
         for (exported_symbol, _) in tcx.exported_symbols(cnum).iter() {
@@ -422,8 +433,13 @@ fn upstream_monomorphizations_provider(
                     if let Some(async_drop_in_place_fn_def_id) = async_drop_in_place_fn_def_id {
                         (async_drop_in_place_fn_def_id, tcx.mk_args(&[ty.into()]))
                     } else {
-                        // `drop_in_place` in place does not exist, don't try
-                        // to use it.
+                        continue;
+                    }
+                }
+                ExportedSymbol::AsyncDropGlue(ty) => {
+                    if let Some(poll_fn_def_id) = async_drop_in_place_poll_fn_def_id {
+                        (poll_fn_def_id, tcx.mk_args(&[ty.into()]))
+                    } else {
                         continue;
                     }
                 }
@@ -575,6 +591,13 @@ pub fn symbol_name_for_instance_in_crate<'tcx>(
                 instantiating_crate,
             )
         }
+        ExportedSymbol::AsyncDropGlue(ty) => {
+            rustc_symbol_mangling::symbol_name_for_instance_in_crate(
+                tcx,
+                Instance::resolve_async_drop_in_place_poll(tcx, ty),
+                instantiating_crate,
+            )
+        }
         ExportedSymbol::NoDefId(symbol_name) => symbol_name.to_string(),
     }
 }
@@ -626,6 +649,7 @@ pub fn linking_symbol_name_for_instance_in_crate<'tcx>(
         // AsyncDropGlueCtorShim always use the Rust calling convention and thus follow the
         // target's default symbol decoration scheme.
         ExportedSymbol::AsyncDropGlueCtorShim(..) => None,
+        ExportedSymbol::AsyncDropGlue(..) => None,
         // NoDefId always follow the target's default symbol decoration scheme.
         ExportedSymbol::NoDefId(..) => None,
         // ThreadLocalShim always follow the target's default symbol decoration scheme.

compiler/rustc_const_eval/src/interpret/call.rs

@@ -552,6 +552,8 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
             | ty::InstanceKind::FnPtrAddrShim(..)
             | ty::InstanceKind::ThreadLocalShim(..)
             | ty::InstanceKind::AsyncDropGlueCtorShim(..)
+            | ty::InstanceKind::AsyncDropGlue(..)
+            | ty::InstanceKind::FutureDropPollShim(..)
             | ty::InstanceKind::Item(_) => {
                 // We need MIR for this fn
                 let Some((body, instance)) = M::find_mir_or_eval_fn(

compiler/rustc_hir/src/lang_items.rs

@@ -182,6 +182,7 @@ language_item_table! {
     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);
     AsyncDropInPlacePoll,    sym::async_drop_in_place_poll, async_drop_in_place_poll_fn, Target::Closure,  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);

compiler/rustc_middle/src/middle/exported_symbols.rs

@@ -44,6 +44,7 @@ pub enum ExportedSymbol<'tcx> {
     Generic(DefId, GenericArgsRef<'tcx>),
     DropGlue(Ty<'tcx>),
     AsyncDropGlueCtorShim(Ty<'tcx>),
+    AsyncDropGlue(Ty<'tcx>),
     ThreadLocalShim(DefId),
     NoDefId(ty::SymbolName<'tcx>),
 }
@@ -63,6 +64,9 @@ impl<'tcx> ExportedSymbol<'tcx> {
             ExportedSymbol::AsyncDropGlueCtorShim(ty) => {
                 tcx.symbol_name(ty::Instance::resolve_async_drop_in_place(tcx, ty))
             }
+            ExportedSymbol::AsyncDropGlue(ty) => {
+                tcx.symbol_name(ty::Instance::resolve_async_drop_in_place_poll(tcx, ty))
+            }
             ExportedSymbol::ThreadLocalShim(def_id) => tcx.symbol_name(ty::Instance {
                 def: ty::InstanceKind::ThreadLocalShim(def_id),
                 args: ty::GenericArgs::empty(),

compiler/rustc_middle/src/mir/mod.rs

@@ -205,7 +205,13 @@ pub struct CoroutineInfo<'tcx> {
     /// Coroutine drop glue. This field is populated after the state transform pass.
     pub coroutine_drop: Option<Body<'tcx>>,
 
-    /// The layout of a coroutine. This field is populated after the state transform pass.
+    /// Coroutine async drop glue.
+    pub coroutine_drop_async: Option<Body<'tcx>>,
+
+    /// When coroutine has sync drop, this is async proxy calling `coroutine_drop` sync impl.
+    pub coroutine_drop_proxy_async: Option<Body<'tcx>>,
+
+    /// The layout of a coroutine. Produced by the state transformation.
     pub coroutine_layout: Option<CoroutineLayout<'tcx>>,
 
     /// If this is a coroutine then record the type of source expression that caused this coroutine
@@ -225,6 +231,8 @@ impl<'tcx> CoroutineInfo<'tcx> {
             yield_ty: Some(yield_ty),
             resume_ty: Some(resume_ty),
             coroutine_drop: None,
+            coroutine_drop_async: None,
+            coroutine_drop_proxy_async: None,
             coroutine_layout: None,
         }
     }
@@ -590,6 +598,26 @@ 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_requires_async_drop(&self) -> bool {
+        self.coroutine_drop_async().is_some()
+    }
+
+    #[inline]
+    pub fn future_drop_poll(&self) -> Option<&Body<'tcx>> {
+        self.coroutine.as_ref().and_then(|coroutine| {
+            coroutine
+                .coroutine_drop_async
+                .as_ref()
+                .or(coroutine.coroutine_drop_proxy_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

@@ -71,6 +71,8 @@ impl<'tcx> MonoItem<'tcx> {
                     // statements, plus one for the terminator.
                     InstanceKind::Item(..)
                     | InstanceKind::DropGlue(..)
+                    | InstanceKind::FutureDropPollShim(..)
+                    | InstanceKind::AsyncDropGlue(..)
                     | InstanceKind::AsyncDropGlueCtorShim(..) => {
                         let mir = tcx.instance_mir(instance.def);
                         mir.basic_blocks.iter().map(|bb| bb.statements.len() + 1).sum()
@@ -419,6 +421,8 @@ impl<'tcx> CodegenUnit<'tcx> {
                             | InstanceKind::CloneShim(..)
                             | InstanceKind::ThreadLocalShim(..)
                             | InstanceKind::FnPtrAddrShim(..)
+                            | InstanceKind::AsyncDropGlue(..)
+                            | InstanceKind::FutureDropPollShim(..)
                             | InstanceKind::AsyncDropGlueCtorShim(..) => None,
                         }
                     }

compiler/rustc_middle/src/mir/pretty.rs

@@ -232,9 +232,20 @@ fn dump_path<'tcx>(
             }));
             s
         }
-        ty::InstanceKind::AsyncDropGlueCtorShim(_, Some(ty)) => {
-            // Unfortunately, pretty-printed typed are not very filename-friendly.
-            // We dome some filtering.
+        ty::InstanceKind::AsyncDropGlueCtorShim(_, ty) => {
+            let mut s = ".".to_owned();
+            s.extend(ty.to_string().chars().filter_map(|c| match c {
+                ' ' => None,
+                ':' | '<' | '>' => Some('_'),
+                c => Some(c),
+            }));
+            s
+        }
+        ty::InstanceKind::AsyncDropGlue(_, ty) => {
+            let ty::Coroutine(_, args) = ty.kind() else {
+                bug!();
+            };
+            let ty = args.first().unwrap().expect_ty();
             let mut s = ".".to_owned();
             s.extend(ty.to_string().chars().filter_map(|c| match c {
                 ' ' => None,
@@ -243,6 +254,21 @@ fn dump_path<'tcx>(
             }));
             s
         }
+        ty::InstanceKind::FutureDropPollShim(_, proxy_cor, impl_cor) => {
+            let mut s = ".".to_owned();
+            s.extend(proxy_cor.to_string().chars().filter_map(|c| match c {
+                ' ' => None,
+                ':' | '<' | '>' => Some('_'),
+                c => Some(c),
+            }));
+            s.push_str(".");
+            s.extend(impl_cor.to_string().chars().filter_map(|c| match c {
+                ' ' => None,
+                ':' | '<' | '>' => Some('_'),
+                c => Some(c),
+            }));
+            s
+        }
         _ => String::new(),
     };
 

compiler/rustc_middle/src/mir/visit.rs

@@ -349,17 +349,21 @@ macro_rules! make_mir_visitor {
                             coroutine_closure_def_id: _def_id,
                             receiver_by_ref: _,
                         } |
-                        ty::InstanceKind::AsyncDropGlueCtorShim(_def_id, None) |
                         ty::InstanceKind::DropGlue(_def_id, None) => {}
 
                         ty::InstanceKind::FnPtrShim(_def_id, ty) |
                         ty::InstanceKind::DropGlue(_def_id, Some(ty)) |
                         ty::InstanceKind::CloneShim(_def_id, ty) |
                         ty::InstanceKind::FnPtrAddrShim(_def_id, ty) |
-                        ty::InstanceKind::AsyncDropGlueCtorShim(_def_id, Some(ty)) => {
+                        ty::InstanceKind::AsyncDropGlue(_def_id, ty) |
+                        ty::InstanceKind::AsyncDropGlueCtorShim(_def_id, ty) => {
                             // FIXME(eddyb) use a better `TyContext` here.
                             self.visit_ty($(& $mutability)? *ty, TyContext::Location(location));
                         }
+                        ty::InstanceKind::FutureDropPollShim(_def_id, proxy_ty, impl_ty) => {
+                            self.visit_ty($(& $mutability)? *proxy_ty, TyContext::Location(location));
+                            self.visit_ty($(& $mutability)? *impl_ty, TyContext::Location(location));
+                        }
                     }
                     self.visit_args(callee_args, location);
                 }

compiler/rustc_middle/src/ty/context.rs

@@ -1562,6 +1562,10 @@ impl<'tcx> TyCtxt<'tcx> {
         self.coroutine_kind(def_id).is_some()
     }
 
+    pub fn is_templated_coroutine(self, def_id: DefId) -> bool {
+        Some(def_id) == self.lang_items().async_drop_in_place_poll_fn()
+    }
+
     /// Returns the movability of the coroutine of `def_id`, or panics
     /// if given a `def_id` that is not a coroutine.
     pub fn coroutine_movability(self, def_id: DefId) -> hir::Movability {