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 <dirbaio@dirbaio.net>

Signed-off-by: Xiangfei Ding <dingxiangfei2009@protonmail.ch>

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

@@ -21,6 +21,7 @@
 
 use std::iter;
 
+use rustc_data_structures::unord::UnordSet;
 use rustc_index::bit_set::{BitMatrix, DenseBitSet};
 use rustc_index::{Idx, IndexSlice, IndexVec};
 use tracing::{debug, trace};
@@ -30,6 +31,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> {
@@ -74,6 +86,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 +116,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 +136,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 +160,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;
@@ -168,18 +187,38 @@ pub(super) fn layout<
         valid_range: WrappingRange { start: 0, end: max_discr },
     };
 
-    let promoted_layouts = ineligible_locals.iter().map(|local| local_layouts[local]);
-    prefix_layouts.push(tag_to_layout(tag));
-    prefix_layouts.extend(promoted_layouts);
+    let upvars_in_unresumed: UnordSet<_> =
+        variant_fields[VariantIdx::new(0)].iter().copied().collect();
+    let promoted_layouts = ineligible_locals.iter().filter_map(|local| {
+        if matches!(pack, PackCoroutineLayout::Classic) && upvars_in_unresumed.contains(&local) {
+            // We do not need to promote upvars, they are already in the upvar region
+            None
+        } else {
+            Some(local_layouts[local])
+        }
+    });
+    // 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 } => {
@@ -213,24 +252,75 @@ 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
         .iter_enumerated()
         .map(|(index, variant_fields)| {
+            // Special case: UNRESUMED overlaps with the upvar region of the prefix,
+            // so that moving upvars may eventually become a no-op.
+            let is_unresumed = index.index() == 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 +344,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

@@ -67,7 +67,9 @@ pub use canon_abi::{ArmCall, CanonAbi, InterruptKind, X86Call};
 pub use extern_abi::CVariadicStatus;
 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/lib.rs

@@ -2478,7 +2478,7 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, '_, 'tcx> {
                 // If the local may have been initialized, and it is now currently being
                 // mutated, then it is justified to be annotated with the `mut`
                 // keyword, since the mutation may be a possible reassignment.
-                if is_local_mutation_allowed != LocalMutationIsAllowed::Yes
+                if !matches!(is_local_mutation_allowed, LocalMutationIsAllowed::Yes)
                     && self.is_local_ever_initialized(local, state).is_some()
                 {
                     self.used_mut.insert(local);

compiler/rustc_borrowck/src/path_utils.rs

@@ -3,7 +3,7 @@ use std::ops::ControlFlow;
 use rustc_abi::FieldIdx;
 use rustc_data_structures::graph::dominators::Dominators;
 use rustc_middle::mir::{BasicBlock, Body, Location, Place, PlaceRef, ProjectionElem};
-use rustc_middle::ty::TyCtxt;
+use rustc_middle::ty::{CapturedPlace, TyCtxt};
 use tracing::debug;
 
 use crate::borrow_set::{BorrowData, BorrowSet, TwoPhaseActivation};
@@ -133,7 +133,7 @@ pub(super) fn borrow_of_local_data(place: Place<'_>) -> bool {
 /// of a closure type.
 pub(crate) fn is_upvar_field_projection<'tcx>(
     tcx: TyCtxt<'tcx>,
-    upvars: &[&rustc_middle::ty::CapturedPlace<'tcx>],
+    upvars: &[&CapturedPlace<'tcx>],
     place_ref: PlaceRef<'tcx>,
     body: &Body<'tcx>,
 ) -> Option<FieldIdx> {

compiler/rustc_codegen_cranelift/src/base.rs

@@ -891,6 +891,9 @@ fn codegen_stmt<'tcx>(fx: &mut FunctionCx<'_, '_, 'tcx>, cur_block: Block, stmt:
                             let variant_dest = lval.downcast_variant(fx, variant_index);
                             (variant_index, variant_dest, active_field_index)
                         }
+                        mir::AggregateKind::Coroutine(_def_id, _args) => {
+                            (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/enums/cpp_like.rs

@@ -721,7 +721,6 @@ fn build_union_fields_for_direct_tag_coroutine<'ll, 'tcx>(
 
     let coroutine_layout = cx.tcx.coroutine_layout(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_count = (variant_range.start.as_u32()..variant_range.end.as_u32()).len();
 
@@ -761,7 +760,6 @@ fn build_union_fields_for_direct_tag_coroutine<'ll, 'tcx>(
                 coroutine_type_and_layout,
                 coroutine_type_di_node,
                 coroutine_layout,
-                common_upvar_names,
             );
 
             let span = coroutine_layout.variant_source_info[variant_index].span;

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

@@ -186,9 +186,6 @@ pub(super) fn build_coroutine_di_node<'ll, 'tcx>(
                 )
             };
 
-            let common_upvar_names =
-                cx.tcx.closure_saved_names_of_captured_variables(coroutine_def_id);
-
             // Build variant struct types
             let variant_struct_type_di_nodes: SmallVec<_> = variants
                 .indices()
@@ -216,7 +213,6 @@ pub(super) fn build_coroutine_di_node<'ll, 'tcx>(
                                 coroutine_type_and_layout,
                                 coroutine_type_di_node,
                                 coroutine_layout,
-                                common_upvar_names,
                             ),
                         source_info,
                     }

compiler/rustc_codegen_ssa/src/mir/operand.rs

@@ -934,13 +934,12 @@ impl<'a, 'tcx, V: CodegenObject> OperandValue<V> {
 }
 
 impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
+    #[instrument(level = "debug", skip(self, bx), ret)]
     fn maybe_codegen_consume_direct(
         &mut self,
         bx: &mut Bx,
         place_ref: mir::PlaceRef<'tcx>,
     ) -> Option<OperandRef<'tcx, Bx::Value>> {
-        debug!("maybe_codegen_consume_direct(place_ref={:?})", place_ref);
-
         match self.locals[place_ref.local] {
             LocalRef::Operand(mut o) => {
                 // We only need to handle the projections that
@@ -980,13 +979,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         }
     }
 
+    #[instrument(level = "debug", skip(self, bx), ret)]
     pub fn codegen_consume(
         &mut self,
         bx: &mut Bx,
         place_ref: mir::PlaceRef<'tcx>,
     ) -> OperandRef<'tcx, Bx::Value> {
-        debug!("codegen_consume(place_ref={:?})", place_ref);
-
         let ty = self.monomorphized_place_ty(place_ref);
         let layout = bx.cx().layout_of(ty);
 
@@ -1005,13 +1003,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         bx.load_operand(place)
     }
 
+    #[instrument(level = "debug", skip(self, bx), ret)]
     pub fn codegen_operand(
         &mut self,
         bx: &mut Bx,
         operand: &mir::Operand<'tcx>,
     ) -> OperandRef<'tcx, Bx::Value> {
-        debug!("codegen_operand(operand={:?})", operand);
-
         match *operand {
             mir::Operand::Copy(ref place) | mir::Operand::Move(ref place) => {
                 self.codegen_consume(bx, place.as_ref())