Introduce -Z pack-coroutine-layout flag

compiler/rustc_abi/src/layout.rs

@@ -4,6 +4,7 @@ use std::ops::Deref;
 use std::range::{RangeFrom, RangeInclusive, RangeToInclusive};
 use std::{cmp, iter};
 
+pub use coroutine::PackCoroutineLayout;
 use rustc_hashes::Hash64;
 use rustc_index::Idx;
 use rustc_index::bit_set::BitMatrix;
@@ -248,17 +249,21 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
     >(
         &self,
         local_layouts: &IndexSlice<LocalIdx, F>,
-        prefix_layouts: IndexVec<FieldIdx, F>,
+        relocated_upvars: &IndexSlice<LocalIdx, Option<LocalIdx>>,
+        upvar_layouts: IndexVec<FieldIdx, F>,
         variant_fields: &IndexSlice<VariantIdx, IndexVec<FieldIdx, LocalIdx>>,
         storage_conflicts: &BitMatrix<LocalIdx, LocalIdx>,
+        pack: PackCoroutineLayout,
         tag_to_layout: impl Fn(Scalar) -> F,
     ) -> LayoutCalculatorResult<FieldIdx, VariantIdx, F> {
         coroutine::layout(
             self,
             local_layouts,
-            prefix_layouts,
+            relocated_upvars,
+            upvar_layouts,
             variant_fields,
             storage_conflicts,
+            pack,
             tag_to_layout,
         )
     }

compiler/rustc_abi/src/layout/coroutine.rs

@@ -30,6 +30,17 @@ use crate::{
     StructKind, TagEncoding, VariantLayout, Variants, WrappingRange,
 };
 
+/// This option controls how coroutine saved locals are packed
+/// into the coroutine state data
+#[derive(Debug, Clone, Copy)]
+pub enum PackCoroutineLayout {
+    /// The classic layout where captures are always promoted to coroutine state prefix
+    Classic,
+    /// Captures are first saved into the `UNRESUMED` state and promoted
+    /// when they are used across more than one suspension
+    CapturesOnly,
+}
+
 /// Overlap eligibility and variant assignment for each CoroutineSavedLocal.
 #[derive(Clone, Debug, PartialEq)]
 enum SavedLocalEligibility<VariantIdx, FieldIdx> {
@@ -74,6 +85,7 @@ fn coroutine_saved_local_eligibility<VariantIdx: Idx, FieldIdx: Idx, LocalIdx: I
             }
         }
     }
+    debug!(?ineligible_locals, "after counting variants containing a saved local");
 
     // Next, check every pair of eligible locals to see if they
     // conflict.
@@ -103,6 +115,7 @@ fn coroutine_saved_local_eligibility<VariantIdx: Idx, FieldIdx: Idx, LocalIdx: I
             trace!("removing local {:?} due to conflict with {:?}", remove, other);
         }
     }
+    debug!(?ineligible_locals, "after checking conflicts");
 
     // Count the number of variants in use. If only one of them, then it is
     // impossible to overlap any locals in our layout. In this case it's
@@ -122,6 +135,7 @@ fn coroutine_saved_local_eligibility<VariantIdx: Idx, FieldIdx: Idx, LocalIdx: I
             }
             ineligible_locals.insert_all();
         }
+        debug!(?ineligible_locals, "after checking used variants");
     }
 
     // Write down the order of our locals that will be promoted to the prefix.
@@ -145,20 +159,24 @@ pub(super) fn layout<
 >(
     calc: &super::LayoutCalculator<impl HasDataLayout>,
     local_layouts: &IndexSlice<LocalIdx, F>,
-    mut prefix_layouts: IndexVec<FieldIdx, F>,
+    _relocated_upvars: &IndexSlice<LocalIdx, Option<LocalIdx>>,
+    upvar_layouts: IndexVec<FieldIdx, F>,
     variant_fields: &IndexSlice<VariantIdx, IndexVec<FieldIdx, LocalIdx>>,
     storage_conflicts: &BitMatrix<LocalIdx, LocalIdx>,
+    pack: PackCoroutineLayout,
     tag_to_layout: impl Fn(Scalar) -> F,
 ) -> super::LayoutCalculatorResult<FieldIdx, VariantIdx, F> {
     use SavedLocalEligibility::*;
 
     let (ineligible_locals, assignments) =
         coroutine_saved_local_eligibility(local_layouts.len(), variant_fields, storage_conflicts);
+    debug!(?ineligible_locals);
 
-    // Build a prefix layout, including "promoting" all ineligible
-    // locals as part of the prefix. We compute the layout of all of
-    // these fields at once to get optimal packing.
-    let tag_index = prefix_layouts.next_index();
+    // Build a prefix layout, consisting of only the state tag and, as per request, upvars
+    let tag_index = match pack {
+        PackCoroutineLayout::CapturesOnly => FieldIdx::new(0),
+        PackCoroutineLayout::Classic => upvar_layouts.next_index(),
+    };
 
     // `variant_fields` already accounts for the reserved variants, so no need to add them.
     let max_discr = (variant_fields.len() - 1) as u128;
@@ -169,18 +187,29 @@ pub(super) fn layout<
     };
 
     let promoted_layouts = ineligible_locals.iter().map(|local| local_layouts[local]);
-    prefix_layouts.push(tag_to_layout(tag));
-    prefix_layouts.extend(promoted_layouts);
+    // FIXME: when we introduce more pack scheme, we need to change the prefix layout here
+    let prefix_layouts: IndexVec<_, _> = match pack {
+        PackCoroutineLayout::Classic => {
+            // Classic scheme packs the states as follows
+            // [ <upvars>.. , <state tag>, <promoted ineligibles>] ++ <variant data>
+            // In addition, UNRESUMED overlaps with the <upvars> part
+            upvar_layouts.into_iter().chain([tag_to_layout(tag)]).chain(promoted_layouts).collect()
+        }
+        PackCoroutineLayout::CapturesOnly => {
+            [tag_to_layout(tag)].into_iter().chain(promoted_layouts).collect()
+        }
+    };
+    debug!(?pack, "prefix_layouts={prefix_layouts:#?}");
     let prefix =
         calc.univariant(&prefix_layouts, &ReprOptions::default(), StructKind::AlwaysSized)?;
 
     let (prefix_size, prefix_align) = (prefix.size, prefix.align);
 
-    // Split the prefix layout into the "outer" fields (upvars and
-    // discriminant) and the "promoted" fields. Promoted fields will
-    // get included in each variant that requested them in
-    // CoroutineLayout.
-    debug!("prefix = {:#?}", prefix);
+    // Split the prefix layout into the discriminant and
+    // the "promoted" fields.
+    // Promoted fields will get included in each variant
+    // that requested them in CoroutineLayout.
+    debug!("prefix={prefix:#?}");
     let (outer_fields, promoted_offsets, promoted_memory_index) = match prefix.fields {
         FieldsShape::Arbitrary { mut offsets, in_memory_order } => {
             // "a" (`0..b_start`) and "b" (`b_start..`) correspond to
@@ -209,6 +238,7 @@ pub(super) fn layout<
         _ => unreachable!(),
     };
 
+    // Here we start to compute layout of each state variant
     let mut size = prefix.size;
     let mut align = prefix.align;
     let variants = variant_fields

compiler/rustc_abi/src/lib.rs

@@ -74,7 +74,7 @@ pub use extern_abi::CVariadicStatus;
 pub use extern_abi::{ExternAbi, all_names};
 pub use layout::{FIRST_VARIANT, FieldIdx, LayoutCalculator, LayoutCalculatorError, VariantIdx};
 #[cfg(feature = "nightly")]
-pub use layout::{Layout, TyAbiInterface, TyAndLayout};
+pub use layout::{Layout, PackCoroutineLayout, TyAbiInterface, TyAndLayout};
 pub use wrapping_range::WrappingRange;
 
 #[derive(Clone, Copy, PartialEq, Eq, Default)]

compiler/rustc_middle/src/mir/pretty.rs

@@ -566,8 +566,9 @@ fn write_coroutine_layout<'tcx>(
     w: &mut dyn io::Write,
     options: PrettyPrintMirOptions,
 ) -> io::Result<()> {
-    let CoroutineLayout { field_tys, variant_fields, variant_source_info, storage_conflicts } =
-        layout;
+    let CoroutineLayout {
+        field_tys, variant_fields, variant_source_info, storage_conflicts, ..
+    } = layout;
 
     writeln!(w, "{INDENT}coroutine layout {{")?;
 

compiler/rustc_middle/src/mir/query.rs

@@ -7,6 +7,7 @@ use rustc_errors::ErrorGuaranteed;
 use rustc_index::IndexVec;
 use rustc_index::bit_set::BitMatrix;
 use rustc_macros::{StableHash, TyDecodable, TyEncodable, TypeFoldable, TypeVisitable};
+use rustc_session::config::PackCoroutineLayout;
 use rustc_span::{Span, Symbol};
 
 use super::{ConstValue, SourceInfo};
@@ -19,8 +20,17 @@ rustc_index::newtype_index! {
     pub struct CoroutineSavedLocal {}
 }
 
-#[derive(Clone, Debug, PartialEq, Eq)]
-#[derive(TyEncodable, TyDecodable, StableHash, TypeFoldable, TypeVisitable)]
+#[derive(
+    Clone,
+    Debug,
+    PartialEq,
+    Eq,
+    TyEncodable,
+    TyDecodable,
+    StableHash,
+    TypeFoldable,
+    TypeVisitable
+)]
 pub struct CoroutineSavedTy<'tcx> {
     pub ty: Ty<'tcx>,
     /// Source info corresponding to the local in the original MIR body.
@@ -32,8 +42,7 @@ pub struct CoroutineSavedTy<'tcx> {
 }
 
 /// The layout of coroutine state.
-#[derive(Clone, PartialEq, Eq)]
-#[derive(TyEncodable, TyDecodable, StableHash, TypeFoldable, TypeVisitable)]
+#[derive(Clone, PartialEq, Eq, TyEncodable, TyDecodable, StableHash, TypeFoldable, TypeVisitable)]
 pub struct CoroutineLayout<'tcx> {
     /// The type of every local stored inside the coroutine.
     pub field_tys: IndexVec<CoroutineSavedLocal, CoroutineSavedTy<'tcx>>,
@@ -52,6 +61,29 @@ pub struct CoroutineLayout<'tcx> {
     #[type_foldable(identity)]
     #[type_visitable(ignore)]
     pub storage_conflicts: BitMatrix<CoroutineSavedLocal, CoroutineSavedLocal>,
+
+    /// This map `A -> B` allows later MIR passes, error reporters
+    /// and layout calculator to relate saved locals `A` sourced from upvars
+    /// and locals `B` that upvars are moved into.
+    ///
+    /// For instance, an upvar `_1.0` is assigned saved local `_s12`,
+    /// see notation of [`CoroutineSavedLocal`], in the UNRESUMED state and
+    /// further moved into the internal saved local `_s13`.
+    /// This map, therefore, establishes the mapping from `_s12` to `_s13`,
+    /// so that their memory layout within the coroutine should be overlapped.
+    #[type_foldable(identity)]
+    #[type_visitable(ignore)]
+    pub relocated_upvars: IndexVec<CoroutineSavedLocal, Option<CoroutineSavedLocal>>,
+
+    /// Coroutine layout packing
+    #[type_foldable(identity)]
+    #[type_visitable(ignore)]
+    pub pack: PackCoroutineLayout,
+}
+
+impl<'tcx> CoroutineLayout<'tcx> {
+    /// The initial state of a coroutine
+    pub const UNRESUMED: VariantIdx = VariantIdx::ZERO;
 }
 
 impl Debug for CoroutineLayout<'_> {
@@ -77,6 +109,7 @@ impl Debug for CoroutineLayout<'_> {
                 map.finish()
             })
             .field("storage_conflicts", &self.storage_conflicts)
+            .field("relocated_upvars", &self.relocated_upvars.debug_map_view())
             .finish()
     }
 }
@@ -98,8 +131,19 @@ pub struct ConstQualifs {
 /// order of the category, thereby influencing diagnostic output.
 ///
 /// See also `rustc_const_eval::borrow_check::constraints`.
-#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
-#[derive(TyEncodable, TyDecodable, StableHash, TypeVisitable, TypeFoldable)]
+#[derive(
+    Copy,
+    Clone,
+    Debug,
+    Eq,
+    PartialEq,
+    Hash,
+    TyEncodable,
+    TyDecodable,
+    StableHash,
+    TypeVisitable,
+    TypeFoldable
+)]
 pub enum ConstraintCategory<'tcx> {
     Return(ReturnConstraint),
     Yield,
@@ -156,15 +200,37 @@ pub enum ConstraintCategory<'tcx> {
     SolverRegionConstraint(Span),
 }
 
-#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
-#[derive(TyEncodable, TyDecodable, StableHash, TypeVisitable, TypeFoldable)]
+#[derive(
+    Copy,
+    Clone,
+    Debug,
+    Eq,
+    PartialEq,
+    Hash,
+    TyEncodable,
+    TyDecodable,
+    StableHash,
+    TypeVisitable,
+    TypeFoldable
+)]
 pub enum ReturnConstraint {
     Normal,
     ClosureUpvar(FieldIdx),
 }
 
-#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
-#[derive(TyEncodable, TyDecodable, StableHash, TypeVisitable, TypeFoldable)]
+#[derive(
+    Copy,
+    Clone,
+    Debug,
+    Eq,
+    PartialEq,
+    Hash,
+    TyEncodable,
+    TyDecodable,
+    StableHash,
+    TypeVisitable,
+    TypeFoldable
+)]
 pub enum AnnotationSource {
     Ascription,
     Declaration,

compiler/rustc_middle/src/ty/mod.rs

@@ -109,7 +109,9 @@ pub use self::typeck_results::{
     UserTypeKind,
 };
 use crate::diagnostics::{OpaqueHiddenTypeMismatch, TypeMismatchReason};
-use crate::mir::{Body, CoroutineLayout, CoroutineSavedLocal, MirPhase, SourceInfo};
+use crate::mir::{
+    Body, CoroutineLayout, CoroutineSavedLocal, CoroutineSavedTy, MirPhase, SourceInfo,
+};
 use crate::query::{IntoQueryKey, Providers};
 use crate::ty;
 use crate::ty::codec::{TyDecoder, TyEncoder};
@@ -2003,24 +2005,40 @@ impl<'tcx> TyCtxt<'tcx> {
         args: GenericArgsRef<'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
-            // and doesn't require any locals. Here is an `empty coroutine's layout`
             let arg_cor_ty = args.first().unwrap().expect_ty();
             if arg_cor_ty.is_coroutine() {
+                // Use the actual upvar type from the coroutine args
+                let upvar_tys = args.as_coroutine().upvar_tys();
+                let upvar_ty =
+                    upvar_tys.first().copied().unwrap_or_else(|| Ty::new_mut_ptr(self, arg_cor_ty));
                 let span = self.def_span(def_id);
                 let source_info = SourceInfo::outermost(span);
-                // Even minimal, empty coroutine has 3 states (RESERVED_VARIANTS),
+                let mut field_tys: IndexVec<CoroutineSavedLocal, CoroutineSavedTy<'tcx>> =
+                    IndexVec::new();
+                let upvar_saved_local = field_tys.push(CoroutineSavedTy {
+                    ty: upvar_ty,
+                    source_info,
+                    ignore_for_traits: true,
+                    debuginfo_name: None,
+                });
+                // Even minimal, the trivial coroutine has 3 states (RESERVED_VARIANTS),
                 // so variant_fields and variant_source_info should have 3 elements.
-                let variant_fields: IndexVec<VariantIdx, IndexVec<FieldIdx, CoroutineSavedLocal>> =
-                    iter::repeat(IndexVec::new()).take(CoroutineArgs::RESERVED_VARIANTS).collect();
+                let mut variant_fields: IndexVec<
+                    VariantIdx,
+                    IndexVec<FieldIdx, CoroutineSavedLocal>,
+                > = iter::repeat(IndexVec::new()).take(CoroutineArgs::RESERVED_VARIANTS).collect();
+                variant_fields[VariantIdx::ZERO].push(upvar_saved_local);
                 let variant_source_info: IndexVec<VariantIdx, SourceInfo> =
                     iter::repeat(source_info).take(CoroutineArgs::RESERVED_VARIANTS).collect();
+                let relocated_upvars: IndexVec<CoroutineSavedLocal, Option<CoroutineSavedLocal>> =
+                    IndexVec::from_raw(vec![Some(upvar_saved_local)]);
                 let proxy_layout = CoroutineLayout {
-                    field_tys: [].into(),
+                    field_tys,
                     variant_fields,
                     variant_source_info,
-                    storage_conflicts: BitMatrix::new(0, 0),
+                    storage_conflicts: BitMatrix::new(1, 1),
+                    relocated_upvars,
+                    pack: rustc_session::config::PackCoroutineLayout::No,
                 };
                 return Ok(self.arena.alloc(proxy_layout));
             } else {

compiler/rustc_mir_transform/src/coroutine/layout.rs

@@ -40,6 +40,7 @@ use rustc_mir_dataflow::impls::{
     always_storage_live_locals,
 };
 use rustc_mir_dataflow::{Analysis, Results, ResultsCursor, ResultsVisitor, visit_results};
+use rustc_session::config::PackCoroutineLayout;
 use rustc_span::Span;
 use rustc_span::def_id::{DefId, LocalDefId};
 use rustc_trait_selection::error_reporting::InferCtxtErrorExt;
@@ -433,8 +434,14 @@ pub(super) fn compute_layout<'tcx>(
         tys[saved_local].debuginfo_name.get_or_insert(var.name);
     }
 
-    let layout =
-        CoroutineLayout { field_tys: tys, variant_fields, variant_source_info, storage_conflicts };
+    let layout = CoroutineLayout {
+        field_tys: tys,
+        variant_fields,
+        variant_source_info,
+        storage_conflicts,
+        relocated_upvars: IndexVec::new(),
+        pack: PackCoroutineLayout::No,
+    };
     debug!(?remap);
     debug!(?layout);
     debug!(?storage_liveness);