Auto merge of #151742 - moulins:variant-layout, r=saethlin

Remove redundant information in `rustc_abi::Variants`



Follow-up to #151040; partially addresses #113988.

Replaces the nested `LayoutData` in `Variants::Multiple` by a new, smaller `VariantLayout` struct, and adjust `LayoutData::for_variant`and the layout algorithm in consequence.  
This PR is best reviewed commit-by-commit.

Author: bors

compiler/rustc_abi/src/layout.rs

@@ -11,7 +11,7 @@ use tracing::{debug, trace};
 use crate::{
     AbiAlign, Align, BackendRepr, FieldsShape, HasDataLayout, IndexSlice, IndexVec, Integer,
     LayoutData, Niche, NonZeroUsize, NumScalableVectors, Primitive, ReprOptions, Scalar, Size,
-    StructKind, TagEncoding, TargetDataLayout, Variants, WrappingRange,
+    StructKind, TagEncoding, TargetDataLayout, VariantLayout, Variants, WrappingRange,
 };
 
 mod coroutine;
@@ -638,6 +638,10 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
         }
 
         let calculate_niche_filling_layout = || -> Option<LayoutData<FieldIdx, VariantIdx>> {
+            struct VariantLayoutInfo {
+                align_abi: Align,
+            }
+
             if repr.inhibit_enum_layout_opt() {
                 return None;
             }
@@ -649,18 +653,22 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
             let mut align = dl.aggregate_align;
             let mut max_repr_align = repr.align;
             let mut unadjusted_abi_align = align;
+            let mut combined_seed = repr.field_shuffle_seed;
 
+            let mut variants_info = IndexVec::<VariantIdx, _>::with_capacity(variants.len());
             let mut variant_layouts = variants
-                .iter_enumerated()
-                .map(|(j, v)| {
-                    let mut st = self.univariant(v, repr, StructKind::AlwaysSized).ok()?;
-                    st.variants = Variants::Single { index: j };
+                .iter()
+                .map(|v| {
+                    let st = self.univariant(v, repr, StructKind::AlwaysSized).ok()?;
+
+                    variants_info.push(VariantLayoutInfo { align_abi: st.align.abi });
 
                     align = align.max(st.align.abi);
                     max_repr_align = max_repr_align.max(st.max_repr_align);
                     unadjusted_abi_align = unadjusted_abi_align.max(st.unadjusted_abi_align);
+                    combined_seed = combined_seed.wrapping_add(st.randomization_seed);
 
-                    Some(st)
+                    Some(VariantLayout::from_layout(st))
                 })
                 .collect::<Option<IndexVec<VariantIdx, _>>>()?;
 
@@ -698,23 +706,16 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                 }
 
                 // Determine if it'll fit after the niche.
-                let this_align = layout.align.abi;
+                let this_align = variants_info[i].align_abi;
                 let this_offset = (niche_offset + niche_size).align_to(this_align);
 
                 if this_offset + layout.size > size {
                     return false;
                 }
 
                 // It'll fit, but we need to make some adjustments.
-                match layout.fields {
-                    FieldsShape::Arbitrary { ref mut offsets, .. } => {
-                        for offset in offsets.iter_mut() {
-                            *offset += this_offset;
-                        }
-                    }
-                    FieldsShape::Primitive | FieldsShape::Array { .. } | FieldsShape::Union(..) => {
-                        panic!("Layout of fields should be Arbitrary for variants")
-                    }
+                for offset in layout.field_offsets.iter_mut() {
+                    *offset += this_offset;
                 }
 
                 // It can't be a Scalar or ScalarPair because the offset isn't 0.
@@ -736,7 +737,7 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                 .iter_enumerated()
                 .all(|(i, layout)| i == largest_variant_index || layout.size == Size::ZERO);
             let same_size = size == variant_layouts[largest_variant_index].size;
-            let same_align = align == variant_layouts[largest_variant_index].align.abi;
+            let same_align = align == variants_info[largest_variant_index].align_abi;
 
             let uninhabited = variant_layouts.iter().all(|v| v.is_uninhabited());
             let abi = if same_size && same_align && others_zst {
@@ -759,11 +760,6 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                 BackendRepr::Memory { sized: true }
             };
 
-            let combined_seed = variant_layouts
-                .iter()
-                .map(|v| v.randomization_seed)
-                .fold(repr.field_shuffle_seed, |acc, seed| acc.wrapping_add(seed));
-
             let layout = LayoutData {
                 variants: Variants::Multiple {
                     tag: niche_scalar,
@@ -856,6 +852,7 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
         let mut align = dl.aggregate_align;
         let mut max_repr_align = repr.align;
         let mut unadjusted_abi_align = align;
+        let mut combined_seed = repr.field_shuffle_seed;
 
         let mut size = Size::ZERO;
 
@@ -879,14 +876,13 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
 
         // Create the set of structs that represent each variant.
         let mut layout_variants = variants
-            .iter_enumerated()
-            .map(|(i, field_layouts)| {
-                let mut st = self.univariant(
+            .iter()
+            .map(|field_layouts| {
+                let st = self.univariant(
                     field_layouts,
                     repr,
                     StructKind::Prefixed(min_ity.size(), prefix_align),
                 )?;
-                st.variants = Variants::Single { index: i };
                 // Find the first field we can't move later
                 // to make room for a larger discriminant.
                 for field_idx in st.fields.index_by_increasing_offset() {
@@ -900,7 +896,8 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                 align = align.max(st.align.abi);
                 max_repr_align = max_repr_align.max(st.max_repr_align);
                 unadjusted_abi_align = unadjusted_abi_align.max(st.unadjusted_abi_align);
-                Ok(st)
+                combined_seed = combined_seed.wrapping_add(st.randomization_seed);
+                Ok(VariantLayout::from_layout(st))
             })
             .collect::<Result<IndexVec<VariantIdx, _>, _>>()?;
 
@@ -955,23 +952,16 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
             let old_ity_size = min_ity.size();
             let new_ity_size = ity.size();
             for variant in &mut layout_variants {
-                match variant.fields {
-                    FieldsShape::Arbitrary { ref mut offsets, .. } => {
-                        for i in offsets {
-                            if *i <= old_ity_size {
-                                assert_eq!(*i, old_ity_size);
-                                *i = new_ity_size;
-                            }
-                        }
-                        // We might be making the struct larger.
-                        if variant.size <= old_ity_size {
-                            variant.size = new_ity_size;
-                        }
-                    }
-                    FieldsShape::Primitive | FieldsShape::Array { .. } | FieldsShape::Union(..) => {
-                        panic!("encountered a non-arbitrary layout during enum layout")
+                for i in &mut variant.field_offsets {
+                    if *i <= old_ity_size {
+                        assert_eq!(*i, old_ity_size);
+                        *i = new_ity_size;
                     }
                 }
+                // We might be making the struct larger.
+                if variant.size <= old_ity_size {
+                    variant.size = new_ity_size;
+                }
             }
         }
 
@@ -996,12 +986,10 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
             let mut common_prim = None;
             let mut common_prim_initialized_in_all_variants = true;
             for (field_layouts, layout_variant) in iter::zip(variants, &layout_variants) {
-                let FieldsShape::Arbitrary { ref offsets, .. } = layout_variant.fields else {
-                    panic!("encountered a non-arbitrary layout during enum layout");
-                };
                 // We skip *all* ZST here and later check if we are good in terms of alignment.
                 // This lets us handle some cases involving aligned ZST.
-                let mut fields = iter::zip(field_layouts, offsets).filter(|p| !p.0.is_zst());
+                let mut fields = iter::zip(field_layouts, &layout_variant.field_offsets)
+                    .filter(|p| !p.0.is_zst());
                 let (field, offset) = match (fields.next(), fields.next()) {
                     (None, None) => {
                         common_prim_initialized_in_all_variants = false;
@@ -1096,25 +1084,17 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
             for variant in &mut layout_variants {
                 // We only do this for variants with fields; the others are not accessed anyway.
                 // Also do not overwrite any already existing "clever" ABIs.
-                if variant.fields.count() > 0
-                    && matches!(variant.backend_repr, BackendRepr::Memory { .. })
+                if matches!(variant.backend_repr, BackendRepr::Memory { .. } if variant.has_fields())
                 {
                     variant.backend_repr = abi;
-                    // Also need to bump up the size and alignment, so that the entire value fits
-                    // in here.
+                    // Also need to bump up the size, so that the entire value fits in here.
                     variant.size = cmp::max(variant.size, size);
-                    variant.align.abi = cmp::max(variant.align.abi, align);
                 }
             }
         }
 
         let largest_niche = Niche::from_scalar(dl, Size::ZERO, tag);
 
-        let combined_seed = layout_variants
-            .iter()
-            .map(|v| v.randomization_seed)
-            .fold(repr.field_shuffle_seed, |acc, seed| acc.wrapping_add(seed));
-
         let tagged_layout = LayoutData {
             variants: Variants::Multiple {
                 tag,

compiler/rustc_abi/src/layout/coroutine.rs

@@ -27,7 +27,7 @@ use tracing::{debug, trace};
 
 use crate::{
     BackendRepr, FieldsShape, HasDataLayout, Integer, LayoutData, Primitive, ReprOptions, Scalar,
-    StructKind, TagEncoding, Variants, WrappingRange,
+    StructKind, TagEncoding, VariantLayout, Variants, WrappingRange,
 };
 
 /// Overlap eligibility and variant assignment for each CoroutineSavedLocal.
@@ -230,7 +230,6 @@ pub(super) fn layout<
                 &ReprOptions::default(),
                 StructKind::Prefixed(prefix_size, prefix_align.abi),
             )?;
-            variant.variants = Variants::Single { index };
 
             let FieldsShape::Arbitrary { offsets, in_memory_order } = variant.fields else {
                 unreachable!();
@@ -281,7 +280,7 @@ pub(super) fn layout<
 
             size = size.max(variant.size);
             align = align.max(variant.align);
-            Ok(variant)
+            Ok(VariantLayout::from_layout(variant))
         })
         .collect::<Result<IndexVec<VariantIdx, _>, _>>()?;
 

compiler/rustc_abi/src/layout/simple.rs

@@ -146,4 +146,28 @@ impl<FieldIdx: Idx, VariantIdx: Idx> LayoutData<FieldIdx, VariantIdx> {
             randomization_seed: Hash64::ZERO,
         }
     }
+
+    /// Returns a layout for an inhabited variant.
+    pub fn for_variant(parent: &Self, index: VariantIdx) -> Self {
+        let layout = match &parent.variants {
+            Variants::Multiple { variants, .. } => &variants[index],
+            _ => panic!("Expected multi-variant layout in `Layout::for_variant`"),
+        };
+
+        Self {
+            fields: FieldsShape::Arbitrary {
+                offsets: layout.field_offsets.clone(),
+                in_memory_order: layout.fields_in_memory_order.clone(),
+            },
+            variants: Variants::Single { index },
+            backend_repr: layout.backend_repr,
+            largest_niche: layout.largest_niche,
+            uninhabited: layout.uninhabited,
+            size: layout.size,
+            align: parent.align,
+            max_repr_align: parent.max_repr_align,
+            unadjusted_abi_align: parent.unadjusted_abi_align,
+            randomization_seed: Hash64::ZERO,
+        }
+    }
 }

compiler/rustc_abi/src/lib.rs

@@ -1974,7 +1974,7 @@ pub enum Variants<FieldIdx: Idx, VariantIdx: Idx> {
         tag: Scalar,
         tag_encoding: TagEncoding<VariantIdx>,
         tag_field: FieldIdx,
-        variants: IndexVec<VariantIdx, LayoutData<FieldIdx, VariantIdx>>,
+        variants: IndexVec<VariantIdx, VariantLayout<FieldIdx>>,
     },
 }
 
@@ -2340,3 +2340,40 @@ pub enum AbiFromStrErr {
     /// no "-unwind" variant can be used here
     NoExplicitUnwind,
 }
+
+// NOTE: This struct is generic over the FieldIdx and VariantIdx for rust-analyzer usage.
+#[derive(PartialEq, Eq, Hash, Clone, Debug)]
+#[cfg_attr(feature = "nightly", derive(HashStable_Generic))]
+pub struct VariantLayout<FieldIdx: Idx> {
+    pub size: Size,
+    pub backend_repr: BackendRepr,
+    pub field_offsets: IndexVec<FieldIdx, Size>,
+    fields_in_memory_order: IndexVec<u32, FieldIdx>,
+    largest_niche: Option<Niche>,
+    uninhabited: bool,
+}
+
+impl<FieldIdx: Idx> VariantLayout<FieldIdx> {
+    pub fn from_layout(layout: LayoutData<FieldIdx, impl Idx>) -> Self {
+        let FieldsShape::Arbitrary { offsets, in_memory_order } = layout.fields else {
+            panic!("Layout of fields should be Arbitrary for variants");
+        };
+
+        Self {
+            size: layout.size,
+            backend_repr: layout.backend_repr,
+            field_offsets: offsets,
+            fields_in_memory_order: in_memory_order,
+            largest_niche: layout.largest_niche,
+            uninhabited: layout.uninhabited,
+        }
+    }
+
+    pub fn is_uninhabited(&self) -> bool {
+        self.uninhabited
+    }
+
+    pub fn has_fields(&self) -> bool {
+        self.field_offsets.len() > 0
+    }
+}

compiler/rustc_middle/src/ty/layout.rs

@@ -765,8 +765,8 @@ where
                 tcx.mk_layout(LayoutData::uninhabited_variant(cx, variant_index, fields))
             }
 
-            Variants::Multiple { ref variants, .. } => {
-                cx.tcx().mk_layout(variants[variant_index].clone())
+            Variants::Multiple { .. } => {
+                cx.tcx().mk_layout(LayoutData::for_variant(&this, variant_index))
             }
         };
 

compiler/rustc_public/src/unstable/convert/stable/abi.rs

@@ -216,11 +216,8 @@ impl<'tcx> Stable<'tcx> for rustc_abi::Variants<rustc_abi::FieldIdx, rustc_abi::
                     tag_field: tag_field.stable(tables, cx),
                     variants: variants
                         .iter()
-                        .map(|v| match &v.fields {
-                            rustc_abi::FieldsShape::Arbitrary { offsets, .. } => VariantFields {
-                                offsets: offsets.iter().as_slice().stable(tables, cx),
-                            },
-                            _ => panic!("variant layout should be Arbitrary"),
+                        .map(|v| VariantFields {
+                            offsets: v.field_offsets.iter().as_slice().stable(tables, cx),
                         })
                         .collect(),
                 }

compiler/rustc_ty_utils/src/layout/invariant.rs

@@ -302,29 +302,16 @@ pub(super) fn layout_sanity_check<'tcx>(cx: &LayoutCx<'tcx>, layout: &TyAndLayou
                 }
             }
             for variant in variants.iter() {
-                // No nested "multiple".
-                assert_matches!(variant.variants, Variants::Single { .. });
-                // Variants should have the same or a smaller size as the full thing,
-                // and same for alignment.
+                // Variants should have the same or a smaller size as the full thing.
                 if variant.size > layout.size {
                     bug!(
                         "Type with size {} bytes has variant with size {} bytes: {layout:#?}",
                         layout.size.bytes(),
                         variant.size.bytes(),
                     )
                 }
-                if variant.align.abi > layout.align.abi {
-                    bug!(
-                        "Type with alignment {} bytes has variant with alignment {} bytes: {layout:#?}",
-                        layout.align.bytes(),
-                        variant.align.bytes(),
-                    )
-                }
                 // Skip empty variants.
-                if variant.size == Size::ZERO
-                    || variant.fields.count() == 0
-                    || variant.is_uninhabited()
-                {
+                if variant.size == Size::ZERO || !variant.has_fields() || variant.is_uninhabited() {
                     // These are never actually accessed anyway, so we can skip the coherence check
                     // for them. They also fail that check, since they may have
                     // a different ABI even when the main type is

src/librustdoc/html/render/type_layout.rs

@@ -54,7 +54,7 @@ pub(crate) fn document_type_layout(cx: &Context<'_>, ty_def_id: DefId) -> impl f
                         span_bug!(tcx.def_span(ty_def_id), "not an adt")
                     };
                     let name = adt.variant(variant_idx).name;
-                    let is_unsized = variant_layout.is_unsized();
+                    let is_unsized = variant_layout.backend_repr.is_unsized();
                     let is_uninhabited = variant_layout.is_uninhabited();
                     let size = variant_layout.size.bytes() - tag_size;
                     let type_layout_size = TypeLayoutSize { is_unsized, is_uninhabited, size };