Prereq6 for async drop - templated coroutine processing and layout

Author: azhogin

Diff for: compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/cpp_like.rs

@@ -721,8 +721,7 @@ fn build_union_fields_for_direct_tag_coroutine<'ll, 'tcx>(
         _ => unreachable!(),
     };
 
-    let coroutine_layout =
-        cx.tcx.coroutine_layout(coroutine_def_id, coroutine_args.kind_ty()).unwrap();
+    let coroutine_layout = cx.tcx.coroutine_layout(coroutine_def_id, coroutine_args.args).unwrap();
 
     let common_upvar_names = cx.tcx.closure_saved_names_of_captured_variables(coroutine_def_id);
     let variant_range = coroutine_args.variant_range(coroutine_def_id, cx.tcx);

Diff for: compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/native.rs

@@ -174,10 +174,8 @@ pub(super) fn build_coroutine_di_node<'ll, 'tcx>(
             DIFlags::FlagZero,
         ),
         |cx, coroutine_type_di_node| {
-            let coroutine_layout = cx
-                .tcx
-                .coroutine_layout(coroutine_def_id, coroutine_args.as_coroutine().kind_ty())
-                .unwrap();
+            let coroutine_layout =
+                cx.tcx.coroutine_layout(coroutine_def_id, coroutine_args).unwrap();
 
             let Variants::Multiple { tag_encoding: TagEncoding::Direct, ref variants, .. } =
                 coroutine_type_and_layout.variants

Diff for: compiler/rustc_middle/src/arena.rs

@@ -9,6 +9,7 @@ macro_rules! arena_types {
     ($macro:path) => (
         $macro!([
             [] layout: rustc_abi::LayoutData<rustc_abi::FieldIdx, rustc_abi::VariantIdx>,
+            [] proxy_coroutine_layout: rustc_middle::mir::CoroutineLayout<'tcx>,
             [] fn_abi: rustc_target::callconv::FnAbi<'tcx, rustc_middle::ty::Ty<'tcx>>,
             // AdtDef are interned and compared by address
             [decode] adt_def: rustc_middle::ty::AdtDefData,

Diff for: compiler/rustc_middle/src/query/mod.rs

@@ -1310,7 +1310,11 @@ rustc_queries! {
     /// Generates a MIR body for the shim.
     query mir_shims(key: ty::InstanceKind<'tcx>) -> &'tcx mir::Body<'tcx> {
         arena_cache
-        desc { |tcx| "generating MIR shim for `{}`", tcx.def_path_str(key.def_id()) }
+        desc {
+            |tcx| "generating MIR shim for `{}`, instance={:?}",
+            tcx.def_path_str(key.def_id()),
+            key
+        }
     }
 
     /// The `symbol_name` query provides the symbol name for calling a

Diff for: compiler/rustc_middle/src/ty/layout.rs

@@ -922,23 +922,58 @@ where
                     i,
                 ),
 
-                ty::Coroutine(def_id, args) => match this.variants {
-                    Variants::Empty => unreachable!(),
-                    Variants::Single { index } => TyMaybeWithLayout::Ty(
-                        args.as_coroutine()
-                            .state_tys(def_id, tcx)
-                            .nth(index.as_usize())
-                            .unwrap()
-                            .nth(i)
-                            .unwrap(),
-                    ),
-                    Variants::Multiple { tag, tag_field, .. } => {
-                        if i == tag_field {
-                            return TyMaybeWithLayout::TyAndLayout(tag_layout(tag));
+                ty::Coroutine(def_id, args) => {
+                    // layout of `async_drop_in_place<T>::{closure}` in case,
+                    // when T is a coroutine, contains this internal coroutine's ref
+                    if tcx.is_templated_coroutine(def_id) {
+                        fn find_impl_coroutine<'tcx>(
+                            tcx: TyCtxt<'tcx>,
+                            mut cor_ty: Ty<'tcx>,
+                        ) -> Ty<'tcx> {
+                            let mut ty = cor_ty;
+                            loop {
+                                if let ty::Coroutine(def_id, args) = ty.kind() {
+                                    cor_ty = ty;
+                                    if tcx.is_templated_coroutine(*def_id) {
+                                        ty = args.first().unwrap().expect_ty();
+                                        continue;
+                                    } else {
+                                        return cor_ty;
+                                    }
+                                } else {
+                                    return cor_ty;
+                                }
+                            }
+                        }
+                        let arg_cor_ty = args.first().unwrap().expect_ty();
+                        if arg_cor_ty.is_coroutine() {
+                            assert!(i == 0);
+                            let impl_cor_ty = find_impl_coroutine(tcx, arg_cor_ty);
+                            return TyMaybeWithLayout::Ty(Ty::new_mut_ref(
+                                tcx,
+                                tcx.lifetimes.re_static,
+                                impl_cor_ty,
+                            ));
                         }
-                        TyMaybeWithLayout::Ty(args.as_coroutine().prefix_tys()[i])
                     }
-                },
+                    match this.variants {
+                        Variants::Empty => unreachable!(),
+                        Variants::Single { index } => TyMaybeWithLayout::Ty(
+                            args.as_coroutine()
+                                .state_tys(def_id, tcx)
+                                .nth(index.as_usize())
+                                .unwrap()
+                                .nth(i)
+                                .unwrap(),
+                        ),
+                        Variants::Multiple { tag, tag_field, .. } => {
+                            if i == tag_field {
+                                return TyMaybeWithLayout::TyAndLayout(tag_layout(tag));
+                            }
+                            TyMaybeWithLayout::Ty(args.as_coroutine().prefix_tys()[i])
+                        }
+                    }
+                }
 
                 ty::Tuple(tys) => TyMaybeWithLayout::Ty(tys[i]),
 

Diff for: compiler/rustc_middle/src/ty/mod.rs

@@ -37,6 +37,7 @@ use rustc_hir::LangItem;
 use rustc_hir::def::{CtorKind, CtorOf, DefKind, DocLinkResMap, LifetimeRes, Res};
 use rustc_hir::def_id::{CrateNum, DefId, DefIdMap, LocalDefId, LocalDefIdMap};
 use rustc_index::IndexVec;
+use rustc_index::bit_set::BitMatrix;
 use rustc_macros::{
     Decodable, Encodable, HashStable, TyDecodable, TyEncodable, TypeFoldable, TypeVisitable,
     extension,
@@ -100,7 +101,7 @@ pub use self::visit::{TypeSuperVisitable, TypeVisitable, TypeVisitableExt, TypeV
 use crate::error::{OpaqueHiddenTypeMismatch, TypeMismatchReason};
 use crate::metadata::ModChild;
 use crate::middle::privacy::EffectiveVisibilities;
-use crate::mir::{Body, CoroutineLayout};
+use crate::mir::{Body, CoroutineLayout, CoroutineSavedLocal, CoroutineSavedTy, SourceInfo};
 use crate::query::{IntoQueryParam, Providers};
 use crate::ty;
 pub use crate::ty::diagnostics::*;
@@ -1740,9 +1741,8 @@ impl<'tcx> TyCtxt<'tcx> {
             | ty::InstanceKind::CloneShim(..)
             | ty::InstanceKind::ThreadLocalShim(..)
             | ty::InstanceKind::FnPtrAddrShim(..)
-            | ty::InstanceKind::AsyncDropGlueCtorShim(..) => self.mir_shims(instance),
-            // async drop glue should be processed specifically, as a templated coroutine
-            ty::InstanceKind::AsyncDropGlue(_, _ty) => todo!(),
+            | ty::InstanceKind::AsyncDropGlueCtorShim(..)
+            | ty::InstanceKind::AsyncDropGlue(..) => self.mir_shims(instance),
         }
     }
 
@@ -1853,16 +1853,17 @@ impl<'tcx> TyCtxt<'tcx> {
         self.def_kind(trait_def_id) == DefKind::TraitAlias
     }
 
-    /// Returns layout of a coroutine. Layout might be unavailable if the
+    /// Returns layout of a non-templated coroutine. Layout might be unavailable if the
     /// coroutine is tainted by errors.
     ///
     /// Takes `coroutine_kind` which can be acquired from the `CoroutineArgs::kind_ty`,
     /// e.g. `args.as_coroutine().kind_ty()`.
-    pub fn coroutine_layout(
+    pub fn ordinary_coroutine_layout(
         self,
         def_id: DefId,
         coroutine_kind_ty: Ty<'tcx>,
     ) -> Option<&'tcx CoroutineLayout<'tcx>> {
+        debug_assert_ne!(Some(def_id), self.lang_items().async_drop_in_place_poll_fn());
         let mir = self.optimized_mir(def_id);
         // Regular coroutine
         if coroutine_kind_ty.is_unit() {
@@ -1892,6 +1893,68 @@ impl<'tcx> TyCtxt<'tcx> {
         }
     }
 
+    /// Returns layout of a templated coroutine. Layout might be unavailable if the
+    /// coroutine is tainted by errors. Atm, the only templated coroutine is
+    /// `async_drop_in_place<T>::{closure}` returned from `async fn async_drop_in_place<T>(..)`.
+    pub fn templated_coroutine_layout(self, ty: Ty<'tcx>) -> Option<&'tcx CoroutineLayout<'tcx>> {
+        self.lang_items().async_drop_in_place_poll_fn().and_then(|def_id| {
+            self.mir_shims(InstanceKind::AsyncDropGlue(def_id, ty)).coroutine_layout_raw()
+        })
+    }
+
+    /// Returns layout of a templated (or not) coroutine. Layout might be unavailable if the
+    /// coroutine is tainted by errors.
+    pub fn coroutine_layout(
+        self,
+        def_id: DefId,
+        args: GenericArgsRef<'tcx>,
+    ) -> Option<&'tcx CoroutineLayout<'tcx>> {
+        if Some(def_id) == self.lang_items().async_drop_in_place_poll_fn() {
+            fn find_impl_coroutine<'tcx>(tcx: TyCtxt<'tcx>, mut cor_ty: Ty<'tcx>) -> Ty<'tcx> {
+                let mut ty = cor_ty;
+                loop {
+                    if let ty::Coroutine(def_id, args) = ty.kind() {
+                        cor_ty = ty;
+                        if tcx.is_templated_coroutine(*def_id) {
+                            ty = args.first().unwrap().expect_ty();
+                            continue;
+                        } else {
+                            return cor_ty;
+                        }
+                    } else {
+                        return cor_ty;
+                    }
+                }
+            }
+            // layout of `async_drop_in_place<T>::{closure}` in case,
+            // when T is a coroutine, contains this internal coroutine's ref
+            let arg_cor_ty = args.first().unwrap().expect_ty();
+            if arg_cor_ty.is_coroutine() {
+                let impl_cor_ty = find_impl_coroutine(self, arg_cor_ty);
+                let impl_ref = Ty::new_mut_ref(self, self.lifetimes.re_static, impl_cor_ty);
+                let span = self.def_span(def_id);
+                let source_info = SourceInfo::outermost(span);
+                let proxy_layout = CoroutineLayout {
+                    field_tys: [CoroutineSavedTy {
+                        ty: impl_ref,
+                        source_info,
+                        ignore_for_traits: true,
+                    }]
+                    .into(),
+                    field_names: [None].into(),
+                    variant_fields: [IndexVec::from([CoroutineSavedLocal::ZERO])].into(),
+                    variant_source_info: [source_info].into(),
+                    storage_conflicts: BitMatrix::new(1, 1),
+                };
+                return Some(self.arena.alloc(proxy_layout));
+            } else {
+                self.templated_coroutine_layout(Ty::new_coroutine(self, def_id, args))
+            }
+        } else {
+            self.ordinary_coroutine_layout(def_id, args.as_coroutine().kind_ty())
+        }
+    }
+
     /// Given the `DefId` of an impl, returns the `DefId` of the trait it implements.
     /// If it implements no trait, returns `None`.
     pub fn trait_id_of_impl(self, def_id: DefId) -> Option<DefId> {

Diff for: compiler/rustc_middle/src/ty/sty.rs

@@ -78,8 +78,7 @@ impl<'tcx> ty::CoroutineArgs<TyCtxt<'tcx>> {
     #[inline]
     fn variant_range(&self, def_id: DefId, tcx: TyCtxt<'tcx>) -> Range<VariantIdx> {
         // FIXME requires optimized MIR
-        FIRST_VARIANT
-            ..tcx.coroutine_layout(def_id, tcx.types.unit).unwrap().variant_fields.next_index()
+        FIRST_VARIANT..tcx.coroutine_layout(def_id, self.args).unwrap().variant_fields.next_index()
     }
 
     /// The discriminant for the given variant. Panics if the `variant_index` is
@@ -139,10 +138,14 @@ impl<'tcx> ty::CoroutineArgs<TyCtxt<'tcx>> {
         def_id: DefId,
         tcx: TyCtxt<'tcx>,
     ) -> impl Iterator<Item: Iterator<Item = Ty<'tcx>> + Captures<'tcx>> {
-        let layout = tcx.coroutine_layout(def_id, self.kind_ty()).unwrap();
+        let layout = tcx.coroutine_layout(def_id, self.args).unwrap();
         layout.variant_fields.iter().map(move |variant| {
             variant.iter().map(move |field| {
-                ty::EarlyBinder::bind(layout.field_tys[*field].ty).instantiate(tcx, self.args)
+                if tcx.is_templated_coroutine(def_id) {
+                    layout.field_tys[*field].ty
+                } else {
+                    ty::EarlyBinder::bind(layout.field_tys[*field].ty).instantiate(tcx, self.args)
+                }
             })
         })
     }

Diff for: compiler/rustc_mir_dataflow/src/value_analysis.rs

@@ -407,6 +407,9 @@ impl<'tcx> Map<'tcx> {
             if exclude.contains(local) {
                 continue;
             }
+            if decl.ty.is_templated_coroutine(tcx) {
+                continue;
+            }
 
             // Create a place for the local.
             debug_assert!(self.locals[local].is_none());

Diff for: compiler/rustc_mir_transform/src/known_panics_lint.rs

@@ -888,7 +888,11 @@ impl CanConstProp {
         };
         for (local, val) in cpv.can_const_prop.iter_enumerated_mut() {
             let ty = body.local_decls[local].ty;
-            if ty.is_union() {
+            if ty.is_templated_coroutine(tcx) {
+                // No const propagation for templated coroutine (AsyncDropGlue)
+                *val = ConstPropMode::NoPropagation;
+                continue;
+            } else if ty.is_union() {
                 // Unions are incompatible with the current implementation of
                 // const prop because Rust has no concept of an active
                 // variant of a union

Diff for: compiler/rustc_mir_transform/src/lib.rs

@@ -624,7 +624,7 @@ fn run_runtime_cleanup_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     }
 }
 
-fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+pub(crate) fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     fn o1<T>(x: T) -> WithMinOptLevel<T> {
         WithMinOptLevel(1, x)
     }

Diff for: compiler/rustc_mir_transform/src/validate.rs

@@ -748,7 +748,9 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                             // since we may be in the process of computing this MIR in the
                             // first place.
                             let layout = if def_id == self.caller_body.source.def_id() {
-                                self.caller_body.coroutine_layout_raw()
+                                self.caller_body
+                                    .coroutine_layout_raw()
+                                    .or_else(|| self.tcx.coroutine_layout(def_id, args))
                             } else if self.tcx.needs_coroutine_by_move_body_def_id(def_id)
                                 && let ty::ClosureKind::FnOnce =
                                     args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap()
@@ -758,7 +760,7 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                                 // Same if this is the by-move body of a coroutine-closure.
                                 self.caller_body.coroutine_layout_raw()
                             } else {
-                                self.tcx.coroutine_layout(def_id, args.as_coroutine().kind_ty())
+                                self.tcx.coroutine_layout(def_id, args)
                             };
 
                             let Some(layout) = layout else {