From c87892f3ef35834f0ca373c01738fe96c94e2a49 Mon Sep 17 00:00:00 2001 From: Xiangfei Ding Date: Fri, 9 Jan 2026 22:53:40 +0000 Subject: [PATCH 1/2] rustc_index: Convenient debugging view of the IndexMap This small utility provides a familiar key-value view of an `IndexMap` in the debugging output. Co-authored-by: Dario Nieuwenhuis Signed-off-by: Xiangfei Ding --- compiler/rustc_index/src/vec.rs | 35 +++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/compiler/rustc_index/src/vec.rs b/compiler/rustc_index/src/vec.rs index 13f0dda180be9..90eb212e69bf5 100644 --- a/compiler/rustc_index/src/vec.rs +++ b/compiler/rustc_index/src/vec.rs @@ -197,6 +197,11 @@ impl IndexVec { pub fn append(&mut self, other: &mut Self) { self.raw.append(&mut other.raw); } + + #[inline] + pub fn debug_map_view(&self) -> IndexSliceMapView<'_, I, T> { + IndexSliceMapView(self.as_slice()) + } } /// `IndexVec` is often used as a map, so it provides some map-like APIs. @@ -220,14 +225,44 @@ impl IndexVec> { pub fn contains(&self, index: I) -> bool { self.get(index).and_then(Option::as_ref).is_some() } + + #[inline] + pub fn debug_map_view_compact(&self) -> IndexSliceMapViewCompact<'_, I, T> { + IndexSliceMapViewCompact(self.as_slice()) + } } +pub struct IndexSliceMapView<'a, I: Idx, T>(&'a IndexSlice); +pub struct IndexSliceMapViewCompact<'a, I: Idx, T>(&'a IndexSlice>); + impl fmt::Debug for IndexVec { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt::Debug::fmt(&self.raw, fmt) } } +impl<'a, I: Idx, T: fmt::Debug> fmt::Debug for IndexSliceMapView<'a, I, T> { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut entries = fmt.debug_map(); + for (idx, val) in self.0.iter_enumerated() { + entries.entry(&idx, val); + } + entries.finish() + } +} + +impl<'a, I: Idx, T: fmt::Debug> fmt::Debug for IndexSliceMapViewCompact<'a, I, T> { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut entries = fmt.debug_map(); + for (idx, val) in self.0.iter_enumerated() { + if let Some(val) = val { + entries.entry(&idx, val); + } + } + entries.finish() + } +} + impl Deref for IndexVec { type Target = IndexSlice; From ce492c59494ca61320e8bed72ecdcd5f84327f16 Mon Sep 17 00:00:00 2001 From: Xiangfei Ding Date: Fri, 9 Jan 2026 22:53:40 +0000 Subject: [PATCH 2/2] rustc_abi,rustc_session: Introduce the pack layout options We will now give users two options, to go with the traditional coroutine layout as of Dec 2025 or the new proposed compact layout. The compact layout will be documented in the RelocateUpvar MIR pass later. Co-authored-by: Dario Nieuwenhuis Signed-off-by: Xiangfei Ding add rustc-hash dependency to rustc_abi --- Cargo.lock | 1 + compiler/rustc_abi/Cargo.toml | 1 + compiler/rustc_abi/src/layout.rs | 9 +- compiler/rustc_abi/src/layout/coroutine.rs | 157 +++++++++++++----- compiler/rustc_abi/src/lib.rs | 2 +- compiler/rustc_middle/src/mir/pretty.rs | 5 +- compiler/rustc_middle/src/mir/query.rs | 86 ++++++++-- compiler/rustc_middle/src/ty/layout.rs | 17 +- compiler/rustc_middle/src/ty/mod.rs | 36 +++- .../src/coroutine/layout.rs | 11 +- compiler/rustc_session/src/config.rs | 36 +++- compiler/rustc_session/src/options.rs | 15 ++ compiler/rustc_ty_utils/src/layout.rs | 8 + 13 files changed, 309 insertions(+), 75 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f55daf9e01c5f..0019bcc23e0e7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3600,6 +3600,7 @@ dependencies = [ "bitflags", "rand 0.9.3", "rand_xoshiro", + "rustc-hash 2.1.1", "rustc_data_structures", "rustc_error_messages", "rustc_errors", diff --git a/compiler/rustc_abi/Cargo.toml b/compiler/rustc_abi/Cargo.toml index 13e0bd8703d9a..9b74400ff76fe 100644 --- a/compiler/rustc_abi/Cargo.toml +++ b/compiler/rustc_abi/Cargo.toml @@ -8,6 +8,7 @@ edition = "2024" bitflags = "2.4.1" rand = { version = "0.9.0", default-features = false, optional = true } rand_xoshiro = { version = "0.7.0", optional = true } +rustc-hash = "2.0.0" rustc_data_structures = { path = "../rustc_data_structures", optional = true } rustc_error_messages = { path = "../rustc_error_messages", optional = true } rustc_errors = { path = "../rustc_errors", optional = true } diff --git a/compiler/rustc_abi/src/layout.rs b/compiler/rustc_abi/src/layout.rs index 2218780092287..f51291aceed21 100644 --- a/compiler/rustc_abi/src/layout.rs +++ b/compiler/rustc_abi/src/layout.rs @@ -4,6 +4,7 @@ use std::ops::Deref; use std::range::RangeInclusive; 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 LayoutCalculator { >( &self, local_layouts: &IndexSlice, - prefix_layouts: IndexVec, + relocated_upvars: &IndexSlice>, + upvar_layouts: IndexVec, variant_fields: &IndexSlice>, storage_conflicts: &BitMatrix, + pack: PackCoroutineLayout, tag_to_layout: impl Fn(Scalar) -> F, ) -> LayoutCalculatorResult { coroutine::layout( self, local_layouts, - prefix_layouts, + relocated_upvars, + upvar_layouts, variant_fields, storage_conflicts, + pack, tag_to_layout, ) } diff --git a/compiler/rustc_abi/src/layout/coroutine.rs b/compiler/rustc_abi/src/layout/coroutine.rs index fd68d06c93829..c428dd0062b7b 100644 --- a/compiler/rustc_abi/src/layout/coroutine.rs +++ b/compiler/rustc_abi/src/layout/coroutine.rs @@ -26,10 +26,21 @@ use rustc_index::{Idx, IndexSlice, IndexVec}; use tracing::{debug, trace}; use crate::{ - BackendRepr, FieldsShape, HasDataLayout, Integer, LayoutData, Primitive, ReprOptions, Scalar, - StructKind, TagEncoding, VariantLayout, Variants, WrappingRange, + Align, BackendRepr, FieldsShape, HasDataLayout, Integer, LayoutData, Primitive, ReprOptions, + Scalar, 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 { @@ -74,6 +85,7 @@ fn coroutine_saved_local_eligibility( calc: &super::LayoutCalculator, local_layouts: &IndexSlice, - mut prefix_layouts: IndexVec, + relocated_upvars: &IndexSlice>, + upvar_layouts: IndexVec, variant_fields: &IndexSlice>, storage_conflicts: &BitMatrix, + pack: PackCoroutineLayout, tag_to_layout: impl Fn(Scalar) -> F, ) -> super::LayoutCalculatorResult { 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,19 +186,39 @@ 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: rustc_hash::FxHashSet<_> = + 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 + // [ .. , , ] ++ + // In addition, UNRESUMED overlaps with the 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); + let prefix_size = prefix.size; - // 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,26 +247,74 @@ 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(), + in_memory_order: (0..tag_index.index()).map(FieldIdx::new).collect(), + }; + let align = prefix.align; + let size = prefix.size; + return Ok(VariantLayout::from_layout(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] + && inner_local != 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::>(), + &fields, &ReprOptions::default(), - StructKind::Prefixed(prefix_size, prefix_align.abi), + StructKind::Prefixed(prefix_size, Align::ONE), )?; let FieldsShape::Arbitrary { offsets, in_memory_order } = variant.fields else { @@ -250,19 +336,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_in_memory_order[memory_index] = i; offset diff --git a/compiler/rustc_abi/src/lib.rs b/compiler/rustc_abi/src/lib.rs index bff4c9bdf47ef..52eedf09b7443 100644 --- a/compiler/rustc_abi/src/lib.rs +++ b/compiler/rustc_abi/src/lib.rs @@ -73,7 +73,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}; #[derive(Clone, Copy, PartialEq, Eq, Default)] #[cfg_attr(feature = "nightly", derive(Encodable_NoContext, Decodable_NoContext, StableHash))] diff --git a/compiler/rustc_middle/src/mir/pretty.rs b/compiler/rustc_middle/src/mir/pretty.rs index 021c1c176d788..6fa3756f226fd 100644 --- a/compiler/rustc_middle/src/mir/pretty.rs +++ b/compiler/rustc_middle/src/mir/pretty.rs @@ -559,8 +559,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 {{")?; diff --git a/compiler/rustc_middle/src/mir/query.rs b/compiler/rustc_middle/src/mir/query.rs index 616b1719359f1..6ccf91a97e1d0 100644 --- a/compiler/rustc_middle/src/mir/query.rs +++ b/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>, @@ -52,6 +61,29 @@ pub struct CoroutineLayout<'tcx> { #[type_foldable(identity)] #[type_visitable(ignore)] pub storage_conflicts: BitMatrix, + + /// 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>, + + /// 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, diff --git a/compiler/rustc_middle/src/ty/layout.rs b/compiler/rustc_middle/src/ty/layout.rs index 3fb35d48513ad..9ce8139d8b00a 100644 --- a/compiler/rustc_middle/src/ty/layout.rs +++ b/compiler/rustc_middle/src/ty/layout.rs @@ -968,14 +968,19 @@ where ty::Coroutine(def_id, args) => match this.variants { Variants::Empty => unreachable!(), - Variants::Single { index } => TyMaybeWithLayout::Ty( - args.as_coroutine() + Variants::Single { index } => { + let mut state_tys = args + .as_coroutine() .state_tys(def_id, tcx) .nth(index.as_usize()) - .unwrap() - .nth(i) - .unwrap(), - ), + .unwrap(); + if let Some(ty) = state_tys.nth(i) { + TyMaybeWithLayout::Ty(ty) + } else { + // Field is not in the variant; it may be an upvar + TyMaybeWithLayout::Ty(args.as_coroutine().upvar_tys()[i]) + } + } Variants::Multiple { tag, tag_field, .. } => { if FieldIdx::from_usize(i) == tag_field { TyMaybeWithLayout::TyAndLayout(tag_layout(tag)) diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs index 9b582eeb2c520..a91c5aa8fa6b9 100644 --- a/compiler/rustc_middle/src/ty/mod.rs +++ b/compiler/rustc_middle/src/ty/mod.rs @@ -121,7 +121,9 @@ pub use self::typeck_results::{ use crate::error::{OpaqueHiddenTypeMismatch, TypeMismatchReason}; use crate::metadata::{AmbigModChild, ModChild}; use crate::middle::privacy::EffectiveVisibilities; -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}; @@ -1996,24 +1998,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::{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> = + 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> = - iter::repeat(IndexVec::new()).take(CoroutineArgs::RESERVED_VARIANTS).collect(); + let mut variant_fields: IndexVec< + VariantIdx, + IndexVec, + > = iter::repeat(IndexVec::new()).take(CoroutineArgs::RESERVED_VARIANTS).collect(); + variant_fields[VariantIdx::ZERO].push(upvar_saved_local); let variant_source_info: IndexVec = iter::repeat(source_info).take(CoroutineArgs::RESERVED_VARIANTS).collect(); + let relocated_upvars: IndexVec> = + 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 { diff --git a/compiler/rustc_mir_transform/src/coroutine/layout.rs b/compiler/rustc_mir_transform/src/coroutine/layout.rs index abf6894e7cde5..2ccbb7f893ee3 100644 --- a/compiler/rustc_mir_transform/src/coroutine/layout.rs +++ b/compiler/rustc_mir_transform/src/coroutine/layout.rs @@ -39,6 +39,7 @@ use rustc_mir_dataflow::impls::{ use rustc_mir_dataflow::{ Analysis, Results, ResultsCursor, ResultsVisitor, visit_reachable_results, }; +use rustc_session::config::PackCoroutineLayout; use rustc_span::Span; use rustc_span::def_id::{DefId, LocalDefId}; use rustc_trait_selection::error_reporting::InferCtxtErrorExt; @@ -441,8 +442,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); diff --git a/compiler/rustc_session/src/config.rs b/compiler/rustc_session/src/config.rs index 37488ebbf1e8f..58118122aaa88 100644 --- a/compiler/rustc_session/src/config.rs +++ b/compiler/rustc_session/src/config.rs @@ -566,8 +566,19 @@ impl SwitchWithOptPath { } } -#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, StableHash)] -#[derive(Encodable, BlobDecodable)] +#[derive( + Copy, + Clone, + Debug, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + StableHash, + Encodable, + BlobDecodable +)] pub enum SymbolManglingVersion { Legacy, V0, @@ -3143,9 +3154,10 @@ pub(crate) mod dep_tracking { CoverageOptions, CrateType, DebugInfo, DebugInfoCompression, ErrorOutputType, FmtDebug, FunctionReturn, InliningThreshold, InstrumentCoverage, InstrumentMcount, InstrumentXRay, LinkerPluginLto, LocationDetail, LtoCli, MirStripDebugInfo, NextSolverConfig, Offload, - OptLevel, OutFileName, OutputType, OutputTypes, PatchableFunctionEntry, PointerAuthOption, - Polonius, ResolveDocLinks, SourceFileHashAlgorithm, SplitDwarfKind, SwitchWithOptPath, - SymbolManglingVersion, WasiExecModel, + OptLevel, OutFileName, OutputType, OutputTypes, PackCoroutineLayout, + PatchableFunctionEntry, PointerAuthOption, Polonius, ResolveDocLinks, + SourceFileHashAlgorithm, SplitDwarfKind, SwitchWithOptPath, SymbolManglingVersion, + WasiExecModel, }; use crate::lint; use crate::utils::NativeLib; @@ -3249,6 +3261,7 @@ pub(crate) mod dep_tracking { Polonius, InliningThreshold, FunctionReturn, + PackCoroutineLayout, Align, CodegenRetagOptions, RustcVersion, @@ -3475,6 +3488,19 @@ pub enum FunctionReturn { ThunkExtern, } +/// Layout optimisation for Coroutines +#[derive(Clone, Copy, PartialEq, Eq, Hash, StableHash, Debug, Default, Decodable, Encodable)] +pub enum PackCoroutineLayout { + /// Keep coroutine captured variables throughout all states + #[default] + No, + + /// Allow coroutine captured variables that are used only once + /// before the first suspension to be freed up for storage + /// in all other suspension states + CapturesOnly, +} + /// Whether extra span comments are included when dumping MIR, via the `-Z mir-include-spans` flag. /// By default, only enabled in the NLL MIR dumps, and disabled in all other passes. #[derive(Clone, Copy, Default, PartialEq, Debug)] diff --git a/compiler/rustc_session/src/options.rs b/compiler/rustc_session/src/options.rs index 5b71c0435185a..f8a02f6698f1b 100644 --- a/compiler/rustc_session/src/options.rs +++ b/compiler/rustc_session/src/options.rs @@ -833,6 +833,7 @@ mod desc { pub(crate) const parse_panic_strategy: &str = "either `unwind`, `abort`, or `immediate-abort`"; pub(crate) const parse_on_broken_pipe: &str = "either `kill`, `error`, or `inherit`"; pub(crate) const parse_patchable_function_entry: &str = "a comma separated list of (prefix_nops,total_nops,section_name), (prefix_nops,total_nops), or (total_nops). Where prefix_nops <= total_nops where 0 < total_nops <= 255 and prefix_nops <= total_nops"; + pub(crate) const parse_pack_coroutine_layout: &str = "either `no` or `captures-only`"; pub(crate) const parse_opt_panic_strategy: &str = parse_panic_strategy; pub(crate) const parse_relro_level: &str = "one of: `full`, `partial`, or `off`"; pub(crate) const parse_sanitizers: &str = "comma separated list of sanitizers: `address`, `cfi`, `dataflow`, `hwaddress`, `kcfi`, `kernel-address`, `kernel-hwaddress`, `leak`, `memory`, `memtag`, `safestack`, `shadow-call-stack`, `thread`, or 'realtime'"; @@ -2075,6 +2076,18 @@ pub mod parse { true } + pub(crate) fn parse_pack_coroutine_layout( + slot: &mut PackCoroutineLayout, + v: Option<&str>, + ) -> bool { + *slot = match v { + Some("no") => PackCoroutineLayout::No, + Some("captures-only") => PackCoroutineLayout::CapturesOnly, + _ => return false, + }; + true + } + pub(crate) fn parse_inlining_threshold(slot: &mut InliningThreshold, v: Option<&str>) -> bool { match v { Some("always" | "yes") => { @@ -2688,6 +2701,8 @@ options! { "behavior of std::io::ErrorKind::BrokenPipe (SIGPIPE)"), osx_rpath_install_name: bool = (false, parse_bool, [TRACKED], "pass `-install_name @rpath/...` to the macOS linker (default: no)"), + pack_coroutine_layout: PackCoroutineLayout = (PackCoroutineLayout::default(), parse_pack_coroutine_layout, [TRACKED], + "set strategy to pack coroutine state layout (default: no)"), packed_bundled_libs: bool = (false, parse_bool, [TRACKED], "change rlib format to store native libraries as archives"), packed_stack: bool = (false, parse_bool, [TRACKED], diff --git a/compiler/rustc_ty_utils/src/layout.rs b/compiler/rustc_ty_utils/src/layout.rs index abec1850502b6..baf8f9e1107e7 100644 --- a/compiler/rustc_ty_utils/src/layout.rs +++ b/compiler/rustc_ty_utils/src/layout.rs @@ -22,6 +22,7 @@ use rustc_middle::ty::{ self, AdtDef, CoroutineArgsExt, EarlyBinder, PseudoCanonicalInput, Ty, TyCtxt, TypeVisitableExt, Unnormalized, }; +use rustc_session::config::PackCoroutineLayout; use rustc_session::{DataTypeKind, FieldInfo, FieldKind, SizeKind, VariantInfo}; use rustc_span::{Symbol, sym}; use tracing::{debug, instrument}; @@ -580,13 +581,20 @@ fn layout_of_uncached<'tcx>( .map(|ty| cx.layout_of(ty)) .try_collect::>()?; + let pack = match info.pack { + PackCoroutineLayout::No => rustc_abi::PackCoroutineLayout::Classic, + PackCoroutineLayout::CapturesOnly => rustc_abi::PackCoroutineLayout::CapturesOnly, + }; + let layout = cx .calc .coroutine( &local_layouts, + &info.relocated_upvars, prefix_layouts, &info.variant_fields, &info.storage_conflicts, + pack, |tag| TyAndLayout { ty: tag.primitive().to_ty(tcx), layout: tcx.mk_layout(LayoutData::scalar(cx, tag)),