new layout

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

Author: dingxiangfei2009

compiler/rustc_abi/src/lib.rs

@@ -1555,8 +1555,10 @@ pub enum Variants<FieldIdx: Idx, VariantIdx: Idx> {
     /// a struct, and they all have space reserved for the tag.
     /// For enums, the tag is the sole field of the layout.
     Multiple {
+        /// Tag definition
         tag: Scalar,
         tag_encoding: TagEncoding<VariantIdx>,
+        /// Index of tag among fields
         tag_field: usize,
         variants: IndexVec<VariantIdx, LayoutData<FieldIdx, VariantIdx>>,
     },
@@ -1600,6 +1602,7 @@ pub enum TagEncoding<VariantIdx: Idx> {
 pub struct Niche {
     pub offset: Size,
     pub value: Primitive,
+    /// A range of valid values with both endpoints being inclusive
     pub valid_range: WrappingRange,
 }
 
@@ -1610,17 +1613,23 @@ impl Niche {
         if niche.available(cx) > 0 { Some(niche) } else { None }
     }
 
+    /// Compute how many values are outside the valid range, available for optimisation through niche filling
     pub fn available<C: HasDataLayout>(&self, cx: &C) -> u128 {
         let Self { value, valid_range: v, .. } = *self;
         let size = value.size(cx);
         assert!(size.bits() <= 128);
         let max_value = size.unsigned_int_max();
 
-        // Find out how many values are outside the valid range.
         let niche = v.end.wrapping_add(1)..v.start;
         niche.end.wrapping_sub(niche.start) & max_value
     }
 
+    /// Try to enlarge the valid value range to include another `count` values,
+    /// so that they can be used for niche-filling optimisation.
+    /// `None` signals impossibility of reservation.
+    /// Otherwise, `Some((start, scalar))` signifies that a reservation is possible,
+    /// the first value in the reservation is `start`, and the new scalar including
+    /// the reserved values is defined in `scalar`.
     pub fn reserve<C: HasDataLayout>(&self, cx: &C, count: u128) -> Option<(u128, Scalar)> {
         assert!(count > 0);
 

compiler/rustc_feature/src/unstable.rs

@@ -453,6 +453,8 @@ declare_features! (
     (internal, contracts_internals, "CURRENT_RUSTC_VERSION", Some(128044)),
     /// Allows coroutines to be cloned.
     (unstable, coroutine_clone, "1.65.0", Some(95360)),
+    /// Allows aggressive merging coroutine saved slots
+    (unstable, coroutine_new_layout, "CURRENT_RUSTC_VERSION", Some(99999)),
     /// Allows defining coroutines.
     (unstable, coroutines, "1.21.0", Some(43122)),
     /// Allows function attribute `#[coverage(on/off)]`, to control coverage

compiler/rustc_middle/src/mir/query.rs

@@ -10,6 +10,7 @@ use rustc_index::bit_set::BitMatrix;
 use rustc_index::{Idx, IndexVec};
 use rustc_macros::{HashStable, TyDecodable, TyEncodable, TypeFoldable, TypeVisitable};
 use rustc_span::{Span, Symbol};
+use rustc_type_ir::data_structures::IndexMap;
 use smallvec::SmallVec;
 
 use super::{ConstValue, SourceInfo};
@@ -57,6 +58,10 @@ pub struct CoroutineLayout<'tcx> {
     #[type_foldable(identity)]
     #[type_visitable(ignore)]
     pub storage_conflicts: BitMatrix<CoroutineSavedLocal, CoroutineSavedLocal>,
+
+    #[type_foldable(identity)]
+    #[type_visitable(ignore)]
+    pub relocated_upvars: IndexMap<CoroutineSavedLocal, CoroutineSavedLocal>,
 }
 
 impl Debug for CoroutineLayout<'_> {

compiler/rustc_mir_transform/src/coroutine.rs

@@ -976,9 +976,17 @@ fn compute_layout<'tcx>(
         suspension_point_at_block,
     } = liveness;
 
+    // We need to later establish the map between upvars in UNRESUMED and locals in other states.
+    let local_upvar_map: UnordMap<_, _> = body
+        .local_upvar_map
+        .iter_enumerated()
+        .filter_map(|(field, local)| local.map(|local| (local, field)))
+        .collect();
+
     // Gather live local types and their indices.
     let mut locals = IndexVec::<CoroutineSavedLocal, _>::new();
     let mut tys = IndexVec::<CoroutineSavedLocal, _>::new();
+    let mut saved_local_upvar_map = UnordMap::default();
     for (saved_local, local) in saved_locals.iter_enumerated() {
         debug!("coroutine saved local {:?} => {:?}", saved_local, local);
 
@@ -1006,6 +1014,10 @@ fn compute_layout<'tcx>(
         debug!(?decl);
 
         tys.push(decl);
+
+        if let Some(&field) = local_upvar_map.get(&local) {
+            saved_local_upvar_map.insert(field, saved_local);
+        }
     }
     // These are the "saved locals" sourced from the UNRESUMED state.
     let upvar_saved_locals: IndexVec<FieldIdx, CoroutineSavedLocal> = upvar_tys
@@ -1058,8 +1070,7 @@ fn compute_layout<'tcx>(
         SourceInfo::outermost(body_span.shrink_to_hi()),
         SourceInfo::outermost(body_span.shrink_to_hi()),
     ]
-    .iter()
-    .copied()
+    .into_iter()
     .collect();
 
     // Build the coroutine variant field list.
@@ -1101,17 +1112,24 @@ fn compute_layout<'tcx>(
             field_names.get_or_insert_with(saved_local, || var.name);
         }
     }
-    for (capture, saved_local) in upvar_infos.iter().zip(upvar_saved_locals) {
+    for (capture, &saved_local) in upvar_infos.iter().zip(&upvar_saved_locals) {
         field_names.get_or_insert_with(saved_local, || capture.var_ident.name);
     }
     debug!(field_names = ?field_names.debug_map_view());
 
+    let relocated_upvars = upvar_saved_locals
+        .iter_enumerated()
+        .filter_map(|(field, &source)| {
+            saved_local_upvar_map.get(&field).map(|&dest| (source, dest))
+        })
+        .collect();
     let layout = CoroutineLayout {
         field_tys: tys,
         field_names,
         variant_fields,
         variant_source_info,
         storage_conflicts,
+        relocated_upvars,
     };
     debug!(?layout);
 

compiler/rustc_mir_transform/src/coroutine/relocate_upvars.rs

@@ -4,24 +4,38 @@
 //! This pass performs the following transformations.
 //! 1. It generates a fresh batch of locals for each captured upvars.
 //!
-//! For each upvar, whether used or not, a fresh local is created with the same type.
+//! For each upvar, whether used or not, a fresh local is created with the same
+//! type.
 //!
-//! 2. It replaces the places pointing into those upvars with places pointing into those locals instead
+//! 2. It replaces the places pointing into those upvars with places pointing
+//! into those locals instead
 //!
-//! Each place that starts with access into the coroutine structure `_1` is replaced with the fresh local as
-//! the base. For instance, `(_1.4 as Some).0` is rewritten into `(_34 as Some).0` when `_34` is the fresh local
+//! Each place that starts with access into the coroutine structure `_1` is
+//! replaced with the fresh local as the base. For instance, `(_1.4 as Some).0`
+//! is rewritten into `(_34 as Some).0` when `_34` is the fresh local
 //! corresponding to the captured upvar stored in `_1.4`.
 //!
 //! 3. It assembles an prologue to replace the current entry block.
 //!
-//! This prologue block transfers every captured upvar into its corresponding fresh local, *via scratch locals*.
-//! The upvars are first completely moved into the scratch locals in batch, and then moved into the destination
-//! locals in batch.
-//! The reason is that it is possible that coroutine layout may change and the source memory location of
-//! an upvar may not necessarily be mapped exactly to the same place as in the `Unresumed` state.
-//! While coroutine layout ensures that the same saved local has stable offsets throughout its lifetime,
-//! technically the upvar in `Unresumed` state and their fresh locals are different saved locals.
-//! This scratch locals re-estabilish safety so that the correct data permutation can take place.
+//! This prologue block transfers every captured upvar into its corresponding
+//! fresh local, *via scratch locals*.
+//! The upvars are first completely moved into the scratch locals in batch,
+//! and then moved into the destination locals in batch.
+//! The reason is that it is possible that coroutine layout may change and the
+//! source memory location of an upvar may not necessarily be mapped exactly to
+//! the same place as in the `Unresumed` state.
+//! While coroutine layout ensures that the same saved local has stable offsets
+//! throughout its lifetime, technically the upvar in `Unresumed` state and
+//! their fresh locals are different saved locals.
+//! This scratch locals re-estabilish safety so that the correct data
+//! permutation can take place.
+//!
+//! By enabling the feature gate `new_coroutine_layout`, the new coroutine
+//! layout calculator enters in effect and further guarantee that the upvars in
+//! the `Unresumed` state will share the same memory offsets as
+//! their corresponding saved locals, if exist.
+//! This policy enables further optimisation opportunities, so that the
+//! copies inserted by this pass will be elided away beyond the codegen phase.
 
 use rustc_abi::FieldIdx;
 use rustc_index::{IndexSlice, IndexVec};

compiler/rustc_span/src/symbol.rs

@@ -704,6 +704,7 @@ symbols! {
         core_panic_macro,
         coroutine,
         coroutine_clone,
+        coroutine_new_layout,
         coroutine_resume,
         coroutine_return,
         coroutine_state,

compiler/rustc_ty_utils/src/layout.rs

@@ -1,3 +1,4 @@
+#![allow(unused)]
 use std::fmt::Debug;
 use std::iter;
 
@@ -31,6 +32,7 @@ use crate::errors::{
     MultipleArrayFieldsSimdType, NonPrimitiveSimdType, OversizedSimdType, ZeroLengthSimdType,
 };
 
+mod coroutine;
 mod invariant;
 
 pub(crate) fn provide(providers: &mut Providers) {
@@ -430,7 +432,13 @@ fn layout_of_uncached<'tcx>(
             tcx.mk_layout(unit)
         }
 
-        ty::Coroutine(def_id, args) => coroutine_layout(cx, ty, def_id, args)?,
+        ty::Coroutine(def_id, args) => {
+            // if tcx.features().coroutine_new_layout() {
+            coroutine::coroutine_layout(cx, ty, def_id, args)?
+            // } else {
+            //     coroutine_layout(cx, ty, def_id, args)?
+            // }
+        }
 
         ty::Closure(_, args) => {
             let tys = args.as_closure().upvar_tys();
@@ -1216,8 +1224,10 @@ fn variant_info_for_coroutine<'tcx>(
         .zip_eq(upvar_names)
         .enumerate()
         .map(|(field_idx, (_, name))| {
-            let field_layout = layout.field(cx, field_idx);
-            let offset = layout.fields.offset(field_idx);
+            // Upvars occupies the Unresumed variant at index zero
+            let variant_layout = layout.for_variant(cx, VariantIdx::ZERO);
+            let field_layout = variant_layout.field(cx, field_idx);
+            let offset = variant_layout.fields.offset(field_idx);
             upvars_size = upvars_size.max(offset + field_layout.size);
             FieldInfo {
                 kind: FieldKind::Upvar,
@@ -1237,12 +1247,11 @@ fn variant_info_for_coroutine<'tcx>(
             let variant_layout = layout.for_variant(cx, variant_idx);
             let mut variant_size = Size::ZERO;
             let fields = variant_def
-                .iter()
-                .enumerate()
+                .iter_enumerated()
                 .map(|(field_idx, local)| {
                     let field_name = coroutine.field_names[*local];
-                    let field_layout = variant_layout.field(cx, field_idx);
-                    let offset = variant_layout.fields.offset(field_idx);
+                    let field_layout = variant_layout.field(cx, field_idx.index());
+                    let offset = variant_layout.fields.offset(field_idx.index());
                     // The struct is as large as the last field's end
                     variant_size = variant_size.max(offset + field_layout.size);
                     FieldInfo {