new layout

Co-authored-by: Dario Nieuwenhuis <dirbaio@dirbaio.net>

Author: dingxiangfei2009

compiler/rustc_abi/src/lib.rs

@@ -1552,8 +1552,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>>,
     },
@@ -1597,6 +1599,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,
 }
 
@@ -1607,17 +1610,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

@@ -447,6 +447,8 @@ declare_features! (
     (unstable, const_try, "1.56.0", Some(74935)),
     /// 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

@@ -11,6 +11,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};
@@ -58,6 +59,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

@@ -975,9 +975,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);
 
@@ -1005,6 +1013,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
@@ -1057,8 +1069,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.
@@ -1100,17 +1111,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_span/src/symbol.rs

@@ -693,6 +693,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) {
@@ -397,7 +399,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();
@@ -1168,8 +1176,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,
@@ -1189,12 +1199,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 {