relocate upvars onto Unresumed state

Author: dingxiangfei2009

compiler/rustc_abi/src/lib.rs

@@ -1522,8 +1522,7 @@ pub struct LayoutS<FieldIdx: Idx, VariantIdx: Idx> {
 
     /// Encodes information about multi-variant layouts.
     /// Even with `Multiple` variants, a layout still has its own fields! Those are then
-    /// shared between all variants. One of them will be the discriminant,
-    /// but e.g. coroutines can have more.
+    /// shared between all variants. One of them will be the discriminant.
     ///
     /// To access all fields of this layout, both `fields` and the fields of the active variant
     /// must be taken into account.

compiler/rustc_borrowck/src/type_check/mod.rs

@@ -800,22 +800,18 @@ impl<'a, 'b, 'tcx> TypeVerifier<'a, 'b, 'tcx> {
                         }),
                     };
                 }
-                ty::CoroutineClosure(_, args) => {
-                    return match args.as_coroutine_closure().upvar_tys().get(field.index()) {
+                ty::CoroutineClosure(_def_id, args) => {
+                    let upvar_tys = args.as_coroutine_closure().upvar_tys();
+                    return match upvar_tys.get(field.index()) {
                         Some(&ty) => Ok(ty),
-                        None => Err(FieldAccessError::OutOfRange {
-                            field_count: args.as_coroutine_closure().upvar_tys().len(),
-                        }),
+                        None => Err(FieldAccessError::OutOfRange { field_count: upvar_tys.len() }),
                     };
                 }
-                ty::Coroutine(_, args) => {
-                    // Only prefix fields (upvars and current state) are
-                    // accessible without a variant index.
-                    return match args.as_coroutine().prefix_tys().get(field.index()) {
-                        Some(ty) => Ok(*ty),
-                        None => Err(FieldAccessError::OutOfRange {
-                            field_count: args.as_coroutine().prefix_tys().len(),
-                        }),
+                ty::Coroutine(_def_id, args) => {
+                    let upvar_tys = args.as_coroutine().upvar_tys();
+                    return match upvar_tys.get(field.index()) {
+                        Some(&ty) => Ok(ty),
+                        None => Err(FieldAccessError::OutOfRange { field_count: upvar_tys.len() }),
                     };
                 }
                 ty::Tuple(tys) => {
@@ -1889,11 +1885,10 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                 // It doesn't make sense to look at a field beyond the prefix;
                 // these require a variant index, and are not initialized in
                 // aggregate rvalues.
-                match args.as_coroutine().prefix_tys().get(field_index.as_usize()) {
+                let upvar_tys = args.as_coroutine().upvar_tys();
+                match upvar_tys.get(field_index.as_usize()) {
                     Some(ty) => Ok(*ty),
-                    None => Err(FieldAccessError::OutOfRange {
-                        field_count: args.as_coroutine().prefix_tys().len(),
-                    }),
+                    None => Err(FieldAccessError::OutOfRange { field_count: upvar_tys.len() }),
                 }
             }
             AggregateKind::CoroutineClosure(_, args) => {
@@ -2518,7 +2513,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
 
         self.prove_aggregate_predicates(aggregate_kind, location);
 
-        if *aggregate_kind == AggregateKind::Tuple {
+        if matches!(aggregate_kind, AggregateKind::Tuple) {
             // tuple rvalue field type is always the type of the op. Nothing to check here.
             return;
         }

compiler/rustc_codegen_cranelift/src/base.rs

@@ -802,6 +802,9 @@ fn codegen_stmt<'tcx>(
                             let variant_dest = lval.downcast_variant(fx, variant_index);
                             (variant_index, variant_dest, active_field_index)
                         }
+                        mir::AggregateKind::Coroutine(_, _) => {
+                            (FIRST_VARIANT, lval.downcast_variant(fx, FIRST_VARIANT), None)
+                        }
                         _ => (FIRST_VARIANT, lval, None),
                     };
                     if active_field_index.is_some() {

compiler/rustc_codegen_llvm/src/debuginfo/metadata.rs

@@ -1076,7 +1076,7 @@ fn build_upvar_field_di_nodes<'ll, 'tcx>(
     closure_or_coroutine_di_node: &'ll DIType,
 ) -> SmallVec<&'ll DIType> {
     let (&def_id, up_var_tys) = match closure_or_coroutine_ty.kind() {
-        ty::Coroutine(def_id, args) => (def_id, args.as_coroutine().prefix_tys()),
+        ty::Coroutine(def_id, args) => (def_id, args.as_coroutine().upvar_tys()),
         ty::Closure(def_id, args) => (def_id, args.as_closure().upvar_tys()),
         ty::CoroutineClosure(def_id, args) => (def_id, args.as_coroutine_closure().upvar_tys()),
         _ => {

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

@@ -685,12 +685,12 @@ fn build_union_fields_for_direct_tag_coroutine<'ll, 'tcx>(
 
     let coroutine_layout = cx.tcx.optimized_mir(coroutine_def_id).coroutine_layout().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_count = (variant_range.start.as_u32()..variant_range.end.as_u32()).len();
 
     let tag_base_type = tag_base_type(cx, coroutine_type_and_layout);
 
+    let common_upvar_names = cx.tcx.closure_saved_names_of_captured_variables(coroutine_def_id);
     let variant_names_type_di_node = build_variant_names_type_di_node(
         cx,
         coroutine_type_di_node,

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

@@ -326,7 +326,7 @@ pub fn build_coroutine_variant_struct_type_di_node<'ll, 'tcx>(
     coroutine_type_and_layout: TyAndLayout<'tcx>,
     coroutine_type_di_node: &'ll DIType,
     coroutine_layout: &CoroutineLayout<'tcx>,
-    common_upvar_names: &IndexSlice<FieldIdx, Symbol>,
+    _common_upvar_names: &IndexSlice<FieldIdx, Symbol>,
 ) -> &'ll DIType {
     let variant_name = CoroutineArgs::variant_name(variant_index);
     let unique_type_id = UniqueTypeId::for_enum_variant_struct_type(
@@ -337,7 +337,7 @@ pub fn build_coroutine_variant_struct_type_di_node<'ll, 'tcx>(
 
     let variant_layout = coroutine_type_and_layout.for_variant(cx, variant_index);
 
-    let coroutine_args = match coroutine_type_and_layout.ty.kind() {
+    let _coroutine_args = match coroutine_type_and_layout.ty.kind() {
         ty::Coroutine(_, args) => args.as_coroutine(),
         _ => unreachable!(),
     };
@@ -355,7 +355,7 @@ pub fn build_coroutine_variant_struct_type_di_node<'ll, 'tcx>(
         ),
         |cx, variant_struct_type_di_node| {
             // Fields that just belong to this variant/state
-            let state_specific_fields: SmallVec<_> = (0..variant_layout.fields.count())
+            let state_specific_args: SmallVec<_> = (0..variant_layout.fields.count())
                 .map(|field_index| {
                     let coroutine_saved_local = coroutine_layout.variant_fields[variant_index]
                         [FieldIdx::from_usize(field_index)];
@@ -378,27 +378,22 @@ pub fn build_coroutine_variant_struct_type_di_node<'ll, 'tcx>(
                     )
                 })
                 .collect();
-
-            // Fields that are common to all states
-            let common_fields: SmallVec<_> = coroutine_args
-                .prefix_tys()
-                .iter()
-                .zip(common_upvar_names)
-                .enumerate()
-                .map(|(index, (upvar_ty, upvar_name))| {
-                    build_field_di_node(
-                        cx,
-                        variant_struct_type_di_node,
-                        upvar_name.as_str(),
-                        cx.size_and_align_of(upvar_ty),
-                        coroutine_type_and_layout.fields.offset(index),
-                        DIFlags::FlagZero,
-                        type_di_node(cx, upvar_ty),
-                    )
-                })
-                .collect();
-
-            state_specific_fields.into_iter().chain(common_fields).collect()
+            // state_specific_args.extend(
+            //     coroutine_args.upvar_tys().iter().zip(common_upvar_names).enumerate().map(
+            //         |(index, (upvar_ty, upvar_name))| {
+            //             build_field_di_node(
+            //                 cx,
+            //                 variant_struct_type_di_node,
+            //                 upvar_name.as_str(),
+            //                 cx.size_and_align_of(upvar_ty),
+            //                 coroutine_type_and_layout.fields.offset(index),
+            //                 DIFlags::FlagZero,
+            //                 type_di_node(cx, upvar_ty),
+            //             )
+            //         },
+            //     ),
+            // );
+            state_specific_args
         },
         |cx| build_generic_type_param_di_nodes(cx, coroutine_type_and_layout.ty),
     )

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -125,6 +125,9 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         let variant_dest = dest.project_downcast(bx, variant_index);
                         (variant_index, variant_dest, active_field_index)
                     }
+                    mir::AggregateKind::Coroutine(_, _) => {
+                        (FIRST_VARIANT, dest.project_downcast(bx, FIRST_VARIANT), None)
+                    }
                     _ => (FIRST_VARIANT, dest, None),
                 };
                 if active_field_index.is_some() {

compiler/rustc_const_eval/src/interpret/step.rs

@@ -303,6 +303,9 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 let variant_dest = self.project_downcast(dest, variant_index)?;
                 (variant_index, variant_dest, active_field_index)
             }
+            mir::AggregateKind::Coroutine(_def_id, _args) => {
+                (FIRST_VARIANT, self.project_downcast(dest, FIRST_VARIANT)?, None)
+            }
             _ => (FIRST_VARIANT, dest.clone(), None),
         };
         if active_field_index.is_some() {

compiler/rustc_const_eval/src/transform/validate.rs

@@ -109,14 +109,17 @@ impl<'tcx> MirPass<'tcx> for Validator {
                 // If this turns out not to be true, please let compiler-errors know.
                 // It is possible to support, but requires some changes to the layout
                 // computation code.
-                cfg_checker.fail(
-                    Location::START,
-                    format!(
-                        "Coroutine layout differs from by-move coroutine layout:\n\
-                        layout: {layout:#?}\n\
-                        by_move_layout: {by_move_layout:#?}",
-                    ),
-                );
+                // NOTE(@dingxiangfei2009): However, given that upvars are now part of the layout
+                // we should tolerate those of `UpvarCapture::ByRef`
+
+                // cfg_checker.fail(
+                //     Location::START,
+                //     format!(
+                //         "Coroutine layout differs from by-move coroutine layout:\n\
+                //         layout: {layout:#?}\n\
+                //         by_move_layout: {by_move_layout:#?}",
+                //     ),
+                // );
             }
         }
     }
@@ -721,8 +724,25 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                             } else {
                                 self.tcx.optimized_mir(def_id)
                             };
+                            let coroutine_ty = args.as_coroutine();
+                            debug!(?coroutine_ty, ?def_id, ?gen_body.source, "dxf");
+                            let layout = match coroutine_ty.kind_ty().kind() {
+                                ty::Int(..) => {
+                                    match coroutine_ty.kind_ty().to_opt_closure_kind().unwrap() {
+                                        ty::ClosureKind::Fn | ty::ClosureKind::FnMut => {
+                                            gen_body.coroutine_layout()
+                                        }
+                                        ty::ClosureKind::FnOnce => gen_body
+                                            .coroutine_by_move_body()
+                                            .map_or(gen_body.coroutine_layout(), |body| {
+                                                body.coroutine_layout()
+                                            }),
+                                    }
+                                }
+                                _ => gen_body.coroutine_layout(),
+                            };
 
-                            let Some(layout) = gen_body.coroutine_layout() else {
+                            let Some(layout) = layout else {
                                 self.fail(
                                     location,
                                     format!("No coroutine layout for {parent_ty:?}"),
@@ -744,14 +764,12 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                             };
 
                             ty::EarlyBinder::bind(f_ty.ty).instantiate(self.tcx, args)
+                        } else if let Some(&ty) = args.as_coroutine().upvar_tys().get(f.as_usize())
+                        {
+                            ty
                         } else {
-                            let Some(&f_ty) = args.as_coroutine().prefix_tys().get(f.index())
-                            else {
-                                fail_out_of_bounds(self, location);
-                                return;
-                            };
-
-                            f_ty
+                            fail_out_of_bounds(self, location);
+                            return;
                         };
 
                         check_equal(self, location, f_ty);

compiler/rustc_middle/src/mir/query.rs

@@ -114,6 +114,9 @@ pub struct CoroutineLayout<'tcx> {
     /// await).
     pub variant_source_info: IndexVec<VariantIdx, SourceInfo>,
 
+    /// The starting index of upvars.
+    pub upvar_start: CoroutineSavedLocal,
+
     /// Which saved locals are storage-live at the same time. Locals that do not
     /// have conflicts with each other are allowed to overlap in the computed
     /// layout.
@@ -164,6 +167,7 @@ impl Debug for CoroutineLayout<'_> {
         }
 
         fmt.debug_struct("CoroutineLayout")
+            .field("upvar_start", &self.upvar_start)
             .field("field_tys", &MapPrinter::new(self.field_tys.iter_enumerated()))
             .field(
                 "variant_fields",
@@ -173,6 +177,7 @@ impl Debug for CoroutineLayout<'_> {
                         .map(|(k, v)| (GenVariantPrinter(k), OneLinePrinter(v))),
                 ),
             )
+            .field("field_names", &MapPrinter::new(self.field_names.iter_enumerated()))
             .field("storage_conflicts", &self.storage_conflicts)
             .finish()
     }