relocate upvars to Unresumed state and make coroutine prefix trivial

Co-authored-by: Dario Nieuwenhuis <[email protected]>

Author: dingxiangfei2009

compiler/rustc_abi/src/layout.rs

@@ -215,15 +215,15 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
     >(
         &self,
         local_layouts: &IndexSlice<LocalIdx, F>,
-        prefix_layouts: IndexVec<FieldIdx, F>,
+        relocated_upvars: &IndexSlice<LocalIdx, Option<LocalIdx>>,
         variant_fields: &IndexSlice<VariantIdx, IndexVec<FieldIdx, LocalIdx>>,
         storage_conflicts: &BitMatrix<LocalIdx, LocalIdx>,
         tag_to_layout: impl Fn(Scalar) -> F,
     ) -> LayoutCalculatorResult<FieldIdx, VariantIdx, F> {
         coroutine::layout(
             self,
             local_layouts,
-            prefix_layouts,
+            relocated_upvars,
             variant_fields,
             storage_conflicts,
             tag_to_layout,

compiler/rustc_abi/src/layout/coroutine.rs

@@ -145,7 +145,7 @@ 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>>,
     variant_fields: &IndexSlice<VariantIdx, IndexVec<FieldIdx, LocalIdx>>,
     storage_conflicts: &BitMatrix<LocalIdx, LocalIdx>,
     tag_to_layout: impl Fn(Scalar) -> F,
@@ -155,10 +155,8 @@ pub(super) fn layout<
     let (ineligible_locals, assignments) =
         coroutine_saved_local_eligibility(local_layouts.len(), variant_fields, storage_conflicts);
 
-    // 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.len();
+    // Build a prefix layout, consisting of only the state tag.
+    let tag_index = 0;
 
     // `variant_fields` already accounts for the reserved variants, so no need to add them.
     let max_discr = (variant_fields.len() - 1) as u128;
@@ -169,17 +167,17 @@ 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);
+    let prefix_layouts: IndexVec<_, _> =
+        [tag_to_layout(tag)].into_iter().chain(promoted_layouts).collect();
     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.
+    // 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, memory_index } => {
@@ -215,22 +213,45 @@ pub(super) fn layout<
 
     let mut size = prefix.size;
     let mut align = prefix.align;
-    let variants = variant_fields
+    let variants: IndexVec<VariantIdx, _> = variant_fields
         .iter_enumerated()
         .map(|(index, variant_fields)| {
+            let is_unresumed = index == VariantIdx::new(0);
+            let mut is_ineligible = IndexVec::from_elem_n(None, variant_fields.len());
+            for (field, &local) in variant_fields.iter_enumerated() {
+                if is_unresumed
+                    && let Some(inner_local) = relocated_upvars[local]
+                    && let Ineligible(Some(promoted_field)) = assignments[inner_local]
+                {
+                    is_ineligible.insert(field, promoted_field);
+                    continue;
+                }
+                match assignments[local] {
+                    Assigned(v) if v == index => {}
+                    Ineligible(Some(promoted_field)) => {
+                        is_ineligible.insert(field, promoted_field);
+                    }
+                    Ineligible(None) => {
+                        panic!("an ineligible local should have been promoted into the prefix")
+                    }
+                    Assigned(_) => {
+                        panic!("an eligible local should have been assigned to exactly one variant")
+                    }
+                    Unassigned => {
+                        panic!("each saved local should have been inspected at least once")
+                    }
+                }
+            }
             // Only include overlap-eligible fields when we compute our variant layout.
-            let variant_only_tys = variant_fields
-                .iter()
-                .filter(|local| match assignments[**local] {
-                    Unassigned => unreachable!(),
-                    Assigned(v) if v == index => true,
-                    Assigned(_) => unreachable!("assignment does not match variant"),
-                    Ineligible(_) => false,
+            let fields: IndexVec<_, _> = variant_fields
+                .iter_enumerated()
+                .filter_map(|(field, &local)| {
+                    if is_ineligible.contains(field) { None } else { Some(local_layouts[local]) }
                 })
-                .map(|local| local_layouts[*local]);
+                .collect();
 
             let mut variant = calc.univariant(
-                &variant_only_tys.collect::<IndexVec<_, _>>(),
+                &fields,
                 &ReprOptions::default(),
                 StructKind::Prefixed(prefix_size, prefix_align.abi),
             )?;
@@ -254,19 +275,14 @@ pub(super) fn layout<
                 IndexVec::from_elem_n(FieldIdx::new(invalid_field_idx), invalid_field_idx);
 
             let mut offsets_and_memory_index = iter::zip(offsets, memory_index);
-            let combined_offsets = variant_fields
+            let combined_offsets = is_ineligible
                 .iter_enumerated()
-                .map(|(i, local)| {
-                    let (offset, memory_index) = match assignments[*local] {
-                        Unassigned => unreachable!(),
-                        Assigned(_) => {
-                            let (offset, memory_index) = offsets_and_memory_index.next().unwrap();
-                            (offset, promoted_memory_index.len() as u32 + memory_index)
-                        }
-                        Ineligible(field_idx) => {
-                            let field_idx = field_idx.unwrap();
-                            (promoted_offsets[field_idx], promoted_memory_index[field_idx])
-                        }
+                .map(|(i, &is_ineligible)| {
+                    let (offset, memory_index) = if let Some(field_idx) = is_ineligible {
+                        (promoted_offsets[field_idx], promoted_memory_index[field_idx])
+                    } else {
+                        let (offset, memory_index) = offsets_and_memory_index.next().unwrap();
+                        (offset, promoted_memory_index.len() as u32 + memory_index)
                     };
                     combined_inverse_memory_index[memory_index] = i;
                     offset
@@ -287,7 +303,7 @@ pub(super) fn layout<
             align = align.max(variant.align);
             Ok(variant)
         })
-        .collect::<Result<IndexVec<VariantIdx, _>, _>>()?;
+        .try_collect()?;
 
     size = size.align_to(align.abi);