Auto merge of #151742 - moulins:variant-layout, r=<try>

Remove redundant information in `rustc_abi::Variants`

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, Primitive, ReprOptions, Scalar, Size, StructKind, TagEncoding,
-    TargetDataLayout, Variants, WrappingRange,
+    TargetDataLayout, VariantLayout, Variants, WrappingRange,
 };
 
 mod coroutine;
@@ -575,6 +575,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;
             }
@@ -586,26 +590,35 @@ 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 largest_variant_size = Size::ZERO;
+            let mut largest_variant_index = VariantIdx::new(0);
+            let mut largest_variant_niche = None;
 
+            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 };
+                .map(|(i, v)| {
+                    let st = self.univariant(v, repr, StructKind::AlwaysSized).ok()?;
+
+                    variants_info.push(VariantLayoutInfo { align_abi: st.align.abi });
+
+                    if st.size > largest_variant_size {
+                        largest_variant_index = i;
+                        largest_variant_size = st.size;
+                        largest_variant_niche = st.largest_niche;
+                    }
 
                     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, _>>>()?;
 
-            let largest_variant_index = variant_layouts
-                .iter_enumerated()
-                .max_by_key(|(_i, layout)| layout.size.bytes())
-                .map(|(i, _layout)| i)?;
-
             let all_indices = variants.indices();
             let needs_disc =
                 |index: VariantIdx| index != largest_variant_index && !absent(&variants[index]);
@@ -616,42 +629,33 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                 (niche_variants.end().index() as u128 - niche_variants.start().index() as u128) + 1;
 
             // Use the largest niche in the largest variant.
-            let niche = variant_layouts[largest_variant_index].largest_niche?;
+            let niche = largest_variant_niche?;
             let (niche_start, niche_scalar) = niche.reserve(dl, count)?;
             let niche_offset = niche.offset;
             let niche_size = niche.value.size(dl);
-            let size = variant_layouts[largest_variant_index].size.align_to(align);
+            let size = largest_variant_size.align_to(align);
 
             let all_variants_fit = variant_layouts.iter_enumerated_mut().all(|(i, layout)| {
                 if i == largest_variant_index {
                     return true;
                 }
 
-                layout.largest_niche = None;
-
                 if layout.size <= niche_offset {
                     // This variant will fit before the niche.
                     return true;
                 }
 
                 // 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.
@@ -673,7 +677,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 {
@@ -696,11 +700,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,
@@ -793,6 +792,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;
 
@@ -816,14 +816,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() {
@@ -837,7 +836,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, _>, _>>()?;
 
@@ -892,23 +892,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;
+                }
             }
         }
 
@@ -933,12 +926,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;
@@ -1033,25 +1024,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: None,
+            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

@@ -1875,7 +1875,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>>,
     },
 }
 
@@ -2230,3 +2230,38 @@ 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>,
+    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,
+            uninhabited: layout.uninhabited,
+        }
+    }
+
+    pub fn is_uninhabited(&self) -> bool {
+        self.uninhabited
+    }
+
+    pub fn has_fields(&self) -> bool {
+        self.field_offsets.len() > 0
+    }
+}

compiler/rustc_codegen_llvm/src/debuginfo/metadata/enums/mod.rs

@@ -1,6 +1,6 @@
 use std::borrow::Cow;
 
-use rustc_abi::{FieldIdx, TagEncoding, VariantIdx, Variants};
+use rustc_abi::{FieldIdx, TagEncoding, VariantIdx, Variants, WrappingRange};
 use rustc_codegen_ssa::debuginfo::type_names::{compute_debuginfo_type_name, cpp_like_debuginfo};
 use rustc_codegen_ssa::debuginfo::{tag_base_type, wants_c_like_enum_debuginfo};
 use rustc_codegen_ssa::traits::MiscCodegenMethods;
@@ -399,16 +399,36 @@ fn compute_discriminant_value<'ll, 'tcx>(
         ),
         &Variants::Multiple {
             tag_encoding: TagEncoding::Niche { ref niche_variants, niche_start, untagged_variant },
+            tag_field,
             tag,
             ..
         } => {
             if variant_index == untagged_variant {
-                let valid_range = enum_type_and_layout
-                    .for_variant(cx, variant_index)
-                    .largest_niche
-                    .as_ref()
-                    .unwrap()
-                    .valid_range;
+                let niche_end = niche_start
+                    .wrapping_sub(niche_variants.start().as_u32() as u128)
+                    .wrapping_add(niche_variants.end().as_u32() as u128);
+
+                // Remove discriminant values of the other variants from the largest niche. This assumes
+                // that the largest niche, when it exists, always corresponds to the enum discriminant.
+                let mut valid_range = WrappingRange {
+                    start: niche_end.wrapping_add(1),
+                    end: niche_start.wrapping_sub(1),
+                };
+                if let Some(niche) = &enum_type_and_layout.largest_niche {
+                    assert_eq!(enum_type_and_layout.fields.offset(tag_field.into()), niche.offset);
+
+                    if niche.valid_range.start == niche_start {
+                        valid_range.end = niche.valid_range.end;
+                    } else if niche.valid_range.end == niche_end {
+                        valid_range.start = niche.valid_range.start;
+                    } else {
+                        bug!(
+                            "largest niche (range: {:?}) doesn't match with niched tag encoding (range: {:?})",
+                            niche.valid_range,
+                            &WrappingRange { start: niche_start, end: niche_end },
+                        )
+                    }
+                }
 
                 let min = valid_range.start.min(valid_range.end);
                 let min = tag.size(cx).truncate(min);