Return value of coroutine_layout fn changed to Result with LayoutError

Author: azhogin

compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/cpp_like.rs

@@ -721,10 +721,14 @@ fn build_union_fields_for_direct_tag_coroutine<'ll, 'tcx>(
         _ => unreachable!(),
     };
 
-    let coroutine_layout = cx.tcx.coroutine_layout(coroutine_def_id, coroutine_args.args).unwrap();
+    let coroutine_layout = cx
+        .tcx
+        .coroutine_layout(coroutine_type_and_layout.ty, 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);
+    let variant_range =
+        coroutine_args.variant_range(coroutine_type_and_layout.ty, coroutine_def_id, cx.tcx);
     let variant_count = (variant_range.start.as_u32()..variant_range.end.as_u32()).len();
 
     let tag_base_type = tag_base_type(cx.tcx, coroutine_type_and_layout);
@@ -743,7 +747,8 @@ fn build_union_fields_for_direct_tag_coroutine<'ll, 'tcx>(
     );
 
     let discriminants: IndexVec<VariantIdx, DiscrResult> = {
-        let discriminants_iter = coroutine_args.discriminants(coroutine_def_id, cx.tcx);
+        let discriminants_iter =
+            coroutine_args.discriminants(coroutine_type_and_layout.ty, coroutine_def_id, cx.tcx);
         let mut discriminants: IndexVec<VariantIdx, DiscrResult> =
             IndexVec::with_capacity(variant_count);
         for (variant_index, discr) in discriminants_iter {

compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/native.rs

@@ -175,7 +175,7 @@ pub(super) fn build_coroutine_di_node<'ll, 'tcx>(
         ),
         |cx, coroutine_type_di_node| {
             let coroutine_layout =
-                cx.tcx.coroutine_layout(coroutine_def_id, coroutine_args).unwrap();
+                cx.tcx.coroutine_layout(coroutine_type, coroutine_def_id, coroutine_args).unwrap();
 
             let Variants::Multiple { tag_encoding: TagEncoding::Direct, ref variants, .. } =
                 coroutine_type_and_layout.variants

compiler/rustc_const_eval/src/interpret/discriminant.rs

@@ -123,7 +123,8 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
                     }
                     ty::Coroutine(def_id, args) => {
                         let args = args.as_coroutine();
-                        args.discriminants(def_id, *self.tcx).find(|(_, var)| var.val == discr_bits)
+                        args.discriminants(ty, def_id, *self.tcx)
+                            .find(|(_, var)| var.val == discr_bits)
                     }
                     _ => span_bug!(self.cur_span(), "tagged layout for non-adt non-coroutine"),
                 }

compiler/rustc_middle/src/mir/statement.rs

@@ -95,7 +95,7 @@ impl<'tcx> PlaceTy<'tcx> {
                     adt_def.variant(variant_index).fields[f].ty(tcx, args)
                 }
                 ty::Coroutine(def_id, args) => {
-                    let mut variants = args.as_coroutine().state_tys(def_id, tcx);
+                    let mut variants = args.as_coroutine().state_tys(self.ty, def_id, tcx);
                     let Some(mut variant) = variants.nth(variant_index.into()) else {
                         bug!("variant {variant_index:?} of coroutine out of range: {self:?}");
                     };

compiler/rustc_middle/src/ty/layout.rs

@@ -927,7 +927,7 @@ where
                     Variants::Empty => unreachable!(),
                     Variants::Single { index } => TyMaybeWithLayout::Ty(
                         args.as_coroutine()
-                            .state_tys(def_id, tcx)
+                            .state_tys(this.ty, def_id, tcx)
                             .nth(index.as_usize())
                             .unwrap()
                             .nth(i)

compiler/rustc_middle/src/ty/mod.rs

@@ -120,6 +120,7 @@ use crate::ty;
 use crate::ty::codec::{TyDecoder, TyEncoder};
 pub use crate::ty::diagnostics::*;
 use crate::ty::fast_reject::SimplifiedType;
+use crate::ty::layout::LayoutError;
 use crate::ty::util::Discr;
 use crate::ty::walk::TypeWalker;
 
@@ -1877,20 +1878,26 @@ impl<'tcx> TyCtxt<'tcx> {
         self.def_kind(trait_def_id) == DefKind::TraitAlias
     }
 
+    /// Arena-alloc of LayoutError for coroutine layout
+    fn layout_error(self, err: LayoutError<'tcx>) -> &'tcx LayoutError<'tcx> {
+        self.arena.alloc(err)
+    }
+
     /// Returns layout of a non-async-drop 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()`.
     fn ordinary_coroutine_layout(
         self,
+        ty: Ty<'tcx>,
         def_id: DefId,
         coroutine_kind_ty: Ty<'tcx>,
-    ) -> Option<&'tcx CoroutineLayout<'tcx>> {
+    ) -> Result<&'tcx CoroutineLayout<'tcx>, &'tcx LayoutError<'tcx>> {
         let mir = self.optimized_mir(def_id);
         // Regular coroutine
         if coroutine_kind_ty.is_unit() {
-            mir.coroutine_layout_raw()
+            mir.coroutine_layout_raw().ok_or_else(|| self.layout_error(LayoutError::Unknown(ty)))
         } else {
             // If we have a `Coroutine` that comes from an coroutine-closure,
             // then it may be a by-move or by-ref body.
@@ -1904,6 +1911,7 @@ impl<'tcx> TyCtxt<'tcx> {
             // a by-ref coroutine.
             if identity_kind_ty == coroutine_kind_ty {
                 mir.coroutine_layout_raw()
+                    .ok_or_else(|| self.layout_error(LayoutError::Unknown(ty)))
             } else {
                 assert_matches!(coroutine_kind_ty.to_opt_closure_kind(), Some(ClosureKind::FnOnce));
                 assert_matches!(
@@ -1912,6 +1920,7 @@ impl<'tcx> TyCtxt<'tcx> {
                 );
                 self.optimized_mir(self.coroutine_by_move_body_def_id(def_id))
                     .coroutine_layout_raw()
+                    .ok_or_else(|| self.layout_error(LayoutError::Unknown(ty)))
             }
         }
     }
@@ -1921,23 +1930,27 @@ impl<'tcx> TyCtxt<'tcx> {
     /// Layout might be unavailable if the coroutine is tainted by errors.
     fn async_drop_coroutine_layout(
         self,
+        ty: Ty<'tcx>,
         def_id: DefId,
         args: GenericArgsRef<'tcx>,
-    ) -> Option<&'tcx CoroutineLayout<'tcx>> {
+    ) -> Result<&'tcx CoroutineLayout<'tcx>, &'tcx LayoutError<'tcx>> {
         if args[0].has_placeholders() || args[0].has_non_region_param() {
-            return None;
+            return Err(self.layout_error(LayoutError::TooGeneric(ty)));
         }
         let instance = InstanceKind::AsyncDropGlue(def_id, Ty::new_coroutine(self, def_id, args));
-        self.mir_shims(instance).coroutine_layout_raw()
+        self.mir_shims(instance)
+            .coroutine_layout_raw()
+            .ok_or_else(|| self.layout_error(LayoutError::Unknown(ty)))
     }
 
     /// Returns layout of a coroutine. Layout might be unavailable if the
     /// coroutine is tainted by errors.
     pub fn coroutine_layout(
         self,
+        ty: Ty<'tcx>,
         def_id: DefId,
         args: GenericArgsRef<'tcx>,
-    ) -> Option<&'tcx CoroutineLayout<'tcx>> {
+    ) -> Result<&'tcx CoroutineLayout<'tcx>, &'tcx LayoutError<'tcx>> {
         if self.is_async_drop_in_place_coroutine(def_id) {
             // layout of `async_drop_in_place<T>::{closure}` in case,
             // when T is a coroutine, contains this internal coroutine's ptr in upvars
@@ -1959,12 +1972,12 @@ impl<'tcx> TyCtxt<'tcx> {
                     variant_source_info,
                     storage_conflicts: BitMatrix::new(0, 0),
                 };
-                return Some(self.arena.alloc(proxy_layout));
+                return Ok(self.arena.alloc(proxy_layout));
             } else {
-                self.async_drop_coroutine_layout(def_id, args)
+                self.async_drop_coroutine_layout(ty, def_id, args)
             }
         } else {
-            self.ordinary_coroutine_layout(def_id, args.as_coroutine().kind_ty())
+            self.ordinary_coroutine_layout(ty, def_id, args.as_coroutine().kind_ty())
         }
     }
 

compiler/rustc_middle/src/ty/sty.rs

@@ -75,23 +75,25 @@ impl<'tcx> ty::CoroutineArgs<TyCtxt<'tcx>> {
 
     /// The valid variant indices of this coroutine.
     #[inline]
-    fn variant_range(&self, def_id: DefId, tcx: TyCtxt<'tcx>) -> Range<VariantIdx> {
+    fn variant_range(&self, ty: Ty<'tcx>, def_id: DefId, tcx: TyCtxt<'tcx>) -> Range<VariantIdx> {
         // FIXME requires optimized MIR
-        FIRST_VARIANT..tcx.coroutine_layout(def_id, self.args).unwrap().variant_fields.next_index()
+        FIRST_VARIANT
+            ..tcx.coroutine_layout(ty, def_id, self.args).unwrap().variant_fields.next_index()
     }
 
     /// The discriminant for the given variant. Panics if the `variant_index` is
     /// out of range.
     #[inline]
     fn discriminant_for_variant(
         &self,
+        ty: Ty<'tcx>,
         def_id: DefId,
         tcx: TyCtxt<'tcx>,
         variant_index: VariantIdx,
     ) -> Discr<'tcx> {
         // Coroutines don't support explicit discriminant values, so they are
         // the same as the variant index.
-        assert!(self.variant_range(def_id, tcx).contains(&variant_index));
+        assert!(self.variant_range(ty, def_id, tcx).contains(&variant_index));
         Discr { val: variant_index.as_usize() as u128, ty: self.discr_ty(tcx) }
     }
 
@@ -100,10 +102,11 @@ impl<'tcx> ty::CoroutineArgs<TyCtxt<'tcx>> {
     #[inline]
     fn discriminants(
         self,
+        ty: Ty<'tcx>,
         def_id: DefId,
         tcx: TyCtxt<'tcx>,
     ) -> impl Iterator<Item = (VariantIdx, Discr<'tcx>)> {
-        self.variant_range(def_id, tcx).map(move |index| {
+        self.variant_range(ty, def_id, tcx).map(move |index| {
             (index, Discr { val: index.as_usize() as u128, ty: self.discr_ty(tcx) })
         })
     }
@@ -134,10 +137,11 @@ impl<'tcx> ty::CoroutineArgs<TyCtxt<'tcx>> {
     #[inline]
     fn state_tys(
         self,
+        ty: Ty<'tcx>,
         def_id: DefId,
         tcx: TyCtxt<'tcx>,
     ) -> impl Iterator<Item: Iterator<Item = Ty<'tcx>>> {
-        let layout = tcx.coroutine_layout(def_id, self.args).unwrap();
+        let layout = tcx.coroutine_layout(ty, def_id, self.args).unwrap();
         layout.variant_fields.iter().map(move |variant| {
             variant.iter().map(move |field| {
                 if tcx.is_async_drop_in_place_coroutine(def_id) {
@@ -1515,7 +1519,7 @@ impl<'tcx> Ty<'tcx> {
         match self.kind() {
             TyKind::Adt(adt, _) => Some(adt.variant_range()),
             TyKind::Coroutine(def_id, args) => {
-                Some(args.as_coroutine().variant_range(*def_id, tcx))
+                Some(args.as_coroutine().variant_range(self, *def_id, tcx))
             }
             _ => None,
         }
@@ -1535,9 +1539,12 @@ impl<'tcx> Ty<'tcx> {
             TyKind::Adt(adt, _) if adt.is_enum() => {
                 Some(adt.discriminant_for_variant(tcx, variant_index))
             }
-            TyKind::Coroutine(def_id, args) => {
-                Some(args.as_coroutine().discriminant_for_variant(*def_id, tcx, variant_index))
-            }
+            TyKind::Coroutine(def_id, args) => Some(args.as_coroutine().discriminant_for_variant(
+                self,
+                *def_id,
+                tcx,
+                variant_index,
+            )),
             _ => None,
         }
     }

compiler/rustc_mir_transform/src/shim.rs

@@ -554,7 +554,7 @@ fn build_clone_shim<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, self_ty: Ty<'tcx>) -
         ty::Tuple(..) => builder.tuple_like_shim(dest, src, self_ty.tuple_fields()),
         ty::Coroutine(coroutine_def_id, args) => {
             assert_eq!(tcx.coroutine_movability(*coroutine_def_id), hir::Movability::Movable);
-            builder.coroutine_shim(dest, src, *coroutine_def_id, args.as_coroutine())
+            builder.coroutine_shim(dest, src, self_ty, *coroutine_def_id, args.as_coroutine())
         }
         _ => bug!("clone shim for `{:?}` which is not `Copy` and is not an aggregate", self_ty),
     };
@@ -754,6 +754,7 @@ impl<'tcx> CloneShimBuilder<'tcx> {
         &mut self,
         dest: Place<'tcx>,
         src: Place<'tcx>,
+        ty: Ty<'tcx>,
         coroutine_def_id: DefId,
         args: CoroutineArgs<TyCtxt<'tcx>>,
     ) {
@@ -764,8 +765,8 @@ impl<'tcx> CloneShimBuilder<'tcx> {
         let unwind = self.clone_fields(dest, src, switch, unwind, args.upvar_tys());
         let target = self.block(vec![], TerminatorKind::Return, false);
         let unreachable = self.block(vec![], TerminatorKind::Unreachable, false);
-        let mut cases = Vec::with_capacity(args.state_tys(coroutine_def_id, self.tcx).count());
-        for (index, state_tys) in args.state_tys(coroutine_def_id, self.tcx).enumerate() {
+        let mut cases = Vec::with_capacity(args.state_tys(ty, coroutine_def_id, self.tcx).count());
+        for (index, state_tys) in args.state_tys(ty, coroutine_def_id, self.tcx).enumerate() {
             let variant_index = VariantIdx::new(index);
             let dest = self.tcx.mk_place_downcast_unnamed(dest, variant_index);
             let src = self.tcx.mk_place_downcast_unnamed(src, variant_index);

compiler/rustc_mir_transform/src/validate.rs

@@ -750,9 +750,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()
-                                    .or_else(|| self.tcx.coroutine_layout(def_id, args))
+                                self.caller_body.coroutine_layout_raw().or_else(|| {
+                                    self.tcx.coroutine_layout(parent_ty.ty, def_id, args).ok()
+                                })
                             } 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()
@@ -762,7 +762,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)
+                                self.tcx.coroutine_layout(parent_ty.ty, def_id, args).ok()
                             };
 
                             let Some(layout) = layout else {

compiler/rustc_ty_utils/src/layout.rs

@@ -492,9 +492,7 @@ fn layout_of_uncached<'tcx>(
         ty::Coroutine(def_id, args) => {
             use rustc_middle::ty::layout::PrimitiveExt as _;
 
-            let Some(info) = tcx.coroutine_layout(def_id, args) else {
-                return Err(error(cx, LayoutError::Unknown(ty)));
-            };
+            let info = tcx.coroutine_layout(ty, def_id, args)?;
 
             let local_layouts = info
                 .field_tys
@@ -856,7 +854,7 @@ fn variant_info_for_coroutine<'tcx>(
         return (vec![], None);
     };
 
-    let coroutine = cx.tcx().coroutine_layout(def_id, args).unwrap();
+    let coroutine = cx.tcx().coroutine_layout(layout.ty, def_id, args).unwrap();
     let upvar_names = cx.tcx().closure_saved_names_of_captured_variables(def_id);
 
     let mut upvars_size = Size::ZERO;