From 81b48563aefb42d3754104235803fa05d749890c Mon Sep 17 00:00:00 2001 From: Xiangfei Ding Date: Fri, 9 Jan 2026 22:53:40 +0000 Subject: [PATCH] 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 --- compiler/rustc_abi/src/layout.rs | 9 +- compiler/rustc_abi/src/layout/coroutine.rs | 54 +++++++++--- 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/mod.rs | 36 ++++++-- .../src/coroutine/layout.rs | 11 ++- compiler/rustc_session/src/config.rs | 35 ++++++-- compiler/rustc_session/src/options.rs | 15 ++++ compiler/rustc_ty_utils/src/layout.rs | 8 ++ 10 files changed, 218 insertions(+), 43 deletions(-) diff --git a/compiler/rustc_abi/src/layout.rs b/compiler/rustc_abi/src/layout.rs index e8779d6ee6869..8f90aa982b148 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::{RangeFrom, RangeInclusive, RangeToInclusive}; 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..7a0e2ffddd9ca 100644 --- a/compiler/rustc_abi/src/layout/coroutine.rs +++ b/compiler/rustc_abi/src/layout/coroutine.rs @@ -30,6 +30,17 @@ use crate::{ 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; @@ -169,18 +187,29 @@ 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 + // [ .. , , ] ++ + // 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); - // 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,6 +238,7 @@ 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 diff --git a/compiler/rustc_abi/src/lib.rs b/compiler/rustc_abi/src/lib.rs index 7cfb93ca1b86d..edd031f7ddf8c 100644 --- a/compiler/rustc_abi/src/lib.rs +++ b/compiler/rustc_abi/src/lib.rs @@ -74,7 +74,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}; pub use wrapping_range::WrappingRange; #[derive(Clone, Copy, PartialEq, Eq, Default)] diff --git a/compiler/rustc_middle/src/mir/pretty.rs b/compiler/rustc_middle/src/mir/pretty.rs index 7bb9b4ff8c375..e2bc0f51fdea7 100644 --- a/compiler/rustc_middle/src/mir/pretty.rs +++ b/compiler/rustc_middle/src/mir/pretty.rs @@ -566,8 +566,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/mod.rs b/compiler/rustc_middle/src/ty/mod.rs index cc6a8619e1e74..ac4e48b27a4ae 100644 --- a/compiler/rustc_middle/src/ty/mod.rs +++ b/compiler/rustc_middle/src/ty/mod.rs @@ -109,7 +109,9 @@ pub use self::typeck_results::{ UserTypeKind, }; use crate::diagnostics::{OpaqueHiddenTypeMismatch, TypeMismatchReason}; -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}; @@ -2003,24 +2005,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 bf2ec025c6381..42f4085f0dfd6 100644 --- a/compiler/rustc_mir_transform/src/coroutine/layout.rs +++ b/compiler/rustc_mir_transform/src/coroutine/layout.rs @@ -40,6 +40,7 @@ use rustc_mir_dataflow::impls::{ always_storage_live_locals, }; use rustc_mir_dataflow::{Analysis, Results, ResultsCursor, ResultsVisitor, visit_results}; +use rustc_session::config::PackCoroutineLayout; use rustc_span::Span; use rustc_span::def_id::{DefId, LocalDefId}; use rustc_trait_selection::error_reporting::InferCtxtErrorExt; @@ -433,8 +434,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 2ae9dfdc9c2ad..e08c0d29025ec 100644 --- a/compiler/rustc_session/src/config.rs +++ b/compiler/rustc_session/src/config.rs @@ -598,8 +598,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, @@ -3346,9 +3357,9 @@ pub(crate) mod dep_tracking { FunctionReturn, InliningThreshold, InstrumentCoverage, InstrumentMcount, InstrumentMcountOpts, InstrumentXRay, LinkerPluginLto, LocationDetail, LtoCli, MirStripDebugInfo, NextSolverConfig, Offload, OptLevel, OutFileName, OutputType, - OutputTypes, PatchableFunctionEntry, PointerAuthOption, Polonius, ResolveDocLinks, - SourceFileHashAlgorithm, SplitDwarfKind, SwitchWithOptPath, SymbolManglingVersion, - WasiExecModel, + OutputTypes, PackCoroutineLayout, PatchableFunctionEntry, PointerAuthOption, Polonius, + ResolveDocLinks, SourceFileHashAlgorithm, SplitDwarfKind, SwitchWithOptPath, + SymbolManglingVersion, WasiExecModel, }; use crate::lint; use crate::utils::NativeLib; @@ -3453,6 +3464,7 @@ pub(crate) mod dep_tracking { Polonius, InliningThreshold, FunctionReturn, + PackCoroutineLayout, Align, CodegenRetagOptions, RustcVersion, @@ -3687,6 +3699,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 8fd9c4da967dc..757886e7f8bd3 100644 --- a/compiler/rustc_session/src/options.rs +++ b/compiler/rustc_session/src/options.rs @@ -832,6 +832,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'"; @@ -2085,6 +2086,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") => { @@ -2716,6 +2729,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 37ff443c83980..2cd1a10d5dd4b 100644 --- a/compiler/rustc_ty_utils/src/layout.rs +++ b/compiler/rustc_ty_utils/src/layout.rs @@ -24,6 +24,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 rustc_structures::Limit; @@ -585,13 +586,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)),