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_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; 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)),