relocate upvars with saved locals for analysis

... and treat coroutine upvar captures as saved locals as well.

This allows the liveness analysis to determine which captures are truly
saved across a yield point and which are initially used but discarded at
first yield points.

In the event that upvar captures are promoted, most certainly because
a coroutine suspends at least once, the slots in the promotion prefix
shall be reused. This means that the copies emitted in the upvar
relocation MIR pass will eventually elided and eliminated in the codegen
phase, hence no additional runtime cost is realised.

Additional MIR dumps are inserted so that it is easier to inspect the
bodies of async closures, including those that captures the state
by-value.

Debug information is updated to point at the correct location for upvars
in borrow checking errors and final debuginfo.

A language change that this patch enables is now actually reverted,
so that lifetimes on relocated upvars are invariant with the upvars outside
of the coroutine body.
We are deferring the language change to a later discussion.

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

Author: dingxiangfei2009

compiler/rustc_abi/src/layout.rs

@@ -3,6 +3,7 @@ use std::fmt::{self, Write};
 use std::ops::{Bound, Deref};
 use std::{cmp, iter};
 
+pub use coroutine::PackCoroutineLayout;
 use rustc_hashes::Hash64;
 use rustc_index::Idx;
 use rustc_index::bit_set::BitMatrix;
@@ -210,17 +211,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, 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 `UNRESUME` 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> {
@@ -145,20 +156,23 @@ 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);
 
-    // 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,17 +183,28 @@ 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, UNRESUME 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!(?prefix_layouts, ?pack);
     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 } => {
@@ -218,19 +243,67 @@ pub(super) fn layout<
     let variants = variant_fields
         .iter_enumerated()
         .map(|(index, variant_fields)| {
+            let is_unresumed = index == VariantIdx::new(0);
+            if is_unresumed && matches!(pack, PackCoroutineLayout::Classic) {
+                let fields = FieldsShape::Arbitrary {
+                    offsets: (0..tag_index.index()).map(|i| outer_fields.offset(i)).collect(),
+                    memory_index: (0..tag_index.index())
+                        .map(|i| {
+                            (outer_fields.memory_index(i) + promoted_memory_index.len()) as u32
+                        })
+                        .collect(),
+                };
+                let align = prefix.align;
+                let size = prefix.size;
+                return Ok(LayoutData {
+                    fields,
+                    variants: Variants::Single { index },
+                    backend_repr: BackendRepr::Memory { sized: true },
+                    largest_niche: None,
+                    uninhabited: false,
+                    align,
+                    size,
+                    max_repr_align: None,
+                    unadjusted_abi_align: align.abi,
+                    randomization_seed: Default::default(),
+                });
+            }
+            let mut is_ineligible = IndexVec::from_elem_n(None, variant_fields.len());
+            for (field, &local) in variant_fields.iter_enumerated() {
+                if is_unresumed {
+                    if 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 +327,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

compiler/rustc_abi/src/lib.rs

@@ -65,7 +65,9 @@ pub use callconv::{Heterogeneous, HomogeneousAggregate, Reg, RegKind};
 pub use canon_abi::{ArmCall, CanonAbi, InterruptKind, X86Call};
 pub use extern_abi::{ExternAbi, all_names};
 #[cfg(feature = "nightly")]
-pub use layout::{FIRST_VARIANT, FieldIdx, Layout, TyAbiInterface, TyAndLayout, VariantIdx};
+pub use layout::{
+    FIRST_VARIANT, FieldIdx, Layout, PackCoroutineLayout, TyAbiInterface, TyAndLayout, VariantIdx,
+};
 pub use layout::{LayoutCalculator, LayoutCalculatorError};
 
 /// Requirements for a `StableHashingContext` to be used in this crate.

compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs

@@ -422,49 +422,18 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                     Place::ty_from(local, proj_base, self.body, self.infcx.tcx).ty
                 ));
 
-                let captured_place = self.upvars[upvar_index.index()];
-
-                err.span_label(span, format!("cannot {act}"));
-
-                let upvar_hir_id = captured_place.get_root_variable();
-
-                if let Node::Pat(pat) = self.infcx.tcx.hir_node(upvar_hir_id)
-                    && let hir::PatKind::Binding(hir::BindingMode::NONE, _, upvar_ident, _) =
-                        pat.kind
-                {
-                    if upvar_ident.name == kw::SelfLower {
-                        for (_, node) in self.infcx.tcx.hir_parent_iter(upvar_hir_id) {
-                            if let Some(fn_decl) = node.fn_decl() {
-                                if !matches!(
-                                    fn_decl.implicit_self,
-                                    hir::ImplicitSelfKind::RefImm | hir::ImplicitSelfKind::RefMut
-                                ) {
-                                    err.span_suggestion_verbose(
-                                        upvar_ident.span.shrink_to_lo(),
-                                        "consider changing this to be mutable",
-                                        "mut ",
-                                        Applicability::MachineApplicable,
-                                    );
-                                    break;
-                                }
-                            }
-                        }
-                    } else {
-                        err.span_suggestion_verbose(
-                            upvar_ident.span.shrink_to_lo(),
-                            "consider changing this to be mutable",
-                            "mut ",
-                            Applicability::MachineApplicable,
-                        );
-                    }
-                }
+                self.suggest_mutable_upvar(*upvar_index, the_place_err, &mut err, span, act);
+            }
 
-                let tcx = self.infcx.tcx;
-                if let ty::Ref(_, ty, Mutability::Mut) = the_place_err.ty(self.body, tcx).ty.kind()
-                    && let ty::Closure(id, _) = *ty.kind()
-                {
-                    self.show_mutating_upvar(tcx, id.expect_local(), the_place_err, &mut err);
-                }
+            PlaceRef { local, projection: [] }
+                if let Some(upvar_index) = self
+                    .body
+                    .local_upvar_map
+                    .iter_enumerated()
+                    .filter_map(|(field, &local)| local.map(|local| (field, local)))
+                    .find_map(|(field, relocated)| (relocated == local).then_some(field)) =>
+            {
+                self.suggest_mutable_upvar(upvar_index, the_place_err, &mut err, span, act);
             }
 
             // Complete hack to approximate old AST-borrowck diagnostic: if the span starts
@@ -570,6 +539,58 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
         }
     }
 
+    fn suggest_mutable_upvar(
+        &self,
+        upvar_index: FieldIdx,
+        the_place_err: PlaceRef<'tcx>,
+        err: &mut Diag<'infcx>,
+        span: Span,
+        act: &str,
+    ) {
+        let captured_place = self.upvars[upvar_index.index()];
+
+        err.span_label(span, format!("cannot {act}"));
+
+        let upvar_hir_id = captured_place.get_root_variable();
+
+        if let Node::Pat(pat) = self.infcx.tcx.hir_node(upvar_hir_id)
+            && let hir::PatKind::Binding(hir::BindingMode::NONE, _, upvar_ident, _) = pat.kind
+        {
+            if upvar_ident.name == kw::SelfLower {
+                for (_, node) in self.infcx.tcx.hir_parent_iter(upvar_hir_id) {
+                    if let Some(fn_decl) = node.fn_decl() {
+                        if !matches!(
+                            fn_decl.implicit_self,
+                            hir::ImplicitSelfKind::RefImm | hir::ImplicitSelfKind::RefMut
+                        ) {
+                            err.span_suggestion_verbose(
+                                upvar_ident.span.shrink_to_lo(),
+                                "consider changing this to be mutable",
+                                "mut ",
+                                Applicability::MachineApplicable,
+                            );
+                            break;
+                        }
+                    }
+                }
+            } else {
+                err.span_suggestion_verbose(
+                    upvar_ident.span.shrink_to_lo(),
+                    "consider changing this to be mutable",
+                    "mut ",
+                    Applicability::MachineApplicable,
+                );
+            }
+        }
+
+        let tcx = self.infcx.tcx;
+        if let ty::Ref(_, ty, Mutability::Mut) = the_place_err.ty(self.body, tcx).ty.kind()
+            && let ty::Closure(id, _) = *ty.kind()
+        {
+            self.show_mutating_upvar(tcx, id.expect_local(), the_place_err, err);
+        }
+    }
+
     /// Suggest `map[k] = v` => `map.insert(k, v)` and the like.
     fn suggest_map_index_mut_alternatives(&self, ty: Ty<'tcx>, err: &mut Diag<'infcx>, span: Span) {
         let Some(adt) = ty.ty_adt_def() else { return };