Rollup merge of #119012 - workingjubilee:extract-enum-layout-fn, r=b-naber

Extract `layout_of_{struct,enum}` fn

While writing #118974 I noticed it was annoying to navigate a huge, several hundred line function, which handles many subcases, and make confident declarations about what part of the flow of execution the compiler would be in. To help with that, this breaks out `layout_of_struct_or_enum`'s fundamental logic into a pair of functions, one for each case. It changes essentially none of that logic, merely moves it around.

Because "the layout of an ADT" feels like a somewhat nebulous subject, I chose to deliberately avoid any expansions to LayoutCalculator's public API, though such does feel like a possible logical next step. There are, indeed, many logical next steps. I'm not taking any of them here, yet, because this comparatively tiny refactor is a prerequisite for all of them.

compiler/rustc_abi/src/layout.rs

@@ -11,6 +11,24 @@ use crate::{
     Variants, WrappingRange,
 };
 
+// A variant is absent if it's uninhabited and only has ZST fields.
+// Present uninhabited variants only require space for their fields,
+// but *not* an encoding of the discriminant (e.g., a tag value).
+// See issue #49298 for more details on the need to leave space
+// for non-ZST uninhabited data (mostly partial initialization).
+fn absent<'a, FieldIdx, VariantIdx, F>(fields: &IndexSlice<FieldIdx, F>) -> bool
+where
+    FieldIdx: Idx,
+    VariantIdx: Idx,
+    F: Deref<Target = &'a LayoutS<FieldIdx, VariantIdx>> + fmt::Debug,
+{
+    let uninhabited = fields.iter().any(|f| f.abi.is_uninhabited());
+    // We cannot ignore alignment; that might lead us to entirely discard a variant and
+    // produce an enum that is less aligned than it should be!
+    let is_1zst = fields.iter().all(|f| f.is_1zst());
+    uninhabited && is_1zst
+}
+
 pub trait LayoutCalculator {
     type TargetDataLayoutRef: Borrow<TargetDataLayout>;
 
@@ -162,24 +180,6 @@ pub trait LayoutCalculator {
         let dl = self.current_data_layout();
         let dl = dl.borrow();
 
-        let scalar_unit = |value: Primitive| {
-            let size = value.size(dl);
-            assert!(size.bits() <= 128);
-            Scalar::Initialized { value, valid_range: WrappingRange::full(size) }
-        };
-
-        // A variant is absent if it's uninhabited and only has ZST fields.
-        // Present uninhabited variants only require space for their fields,
-        // but *not* an encoding of the discriminant (e.g., a tag value).
-        // See issue #49298 for more details on the need to leave space
-        // for non-ZST uninhabited data (mostly partial initialization).
-        let absent = |fields: &IndexSlice<FieldIdx, F>| {
-            let uninhabited = fields.iter().any(|f| f.abi.is_uninhabited());
-            // We cannot ignore alignment; that might lead us to entirely discard a variant and
-            // produce an enum that is less aligned than it should be!
-            let is_1zst = fields.iter().all(|f| f.is_1zst());
-            uninhabited && is_1zst
-        };
         let (present_first, present_second) = {
             let mut present_variants = variants
                 .iter_enumerated()
@@ -197,12 +197,160 @@ pub trait LayoutCalculator {
             None => VariantIdx::new(0),
         };
 
-        let is_struct = !is_enum ||
+        // take the struct path if it is an actual struct
-                    // Only one variant is present.
+        if !is_enum ||
-                    (present_second.is_none() &&
+            // or for optimizing univariant enums
-                        // Representation optimizations are allowed.
+            (present_second.is_none() && !repr.inhibit_enum_layout_opt())
-                        !repr.inhibit_enum_layout_opt());
+        {
-        if is_struct {
+            layout_of_struct(
+                self,
+                repr,
+                variants,
+                is_enum,
+                is_unsafe_cell,
+                scalar_valid_range,
+                always_sized,
+                dl,
+                present_first,
+            )
+        } else {
+            // At this point, we have handled all unions and
+            // structs. (We have also handled univariant enums
+            // that allow representation optimization.)
+            assert!(is_enum);
+            layout_of_enum(
+                self,
+                repr,
+                variants,
+                discr_range_of_repr,
+                discriminants,
+                dont_niche_optimize_enum,
+                dl,
+            )
+        }
+    }
+
+    fn layout_of_union<
+        'a,
+        FieldIdx: Idx,
+        VariantIdx: Idx,
+        F: Deref<Target = &'a LayoutS<FieldIdx, VariantIdx>> + fmt::Debug,
+    >(
+        &self,
+        repr: &ReprOptions,
+        variants: &IndexSlice<VariantIdx, IndexVec<FieldIdx, F>>,
+    ) -> Option<LayoutS<FieldIdx, VariantIdx>> {
+        let dl = self.current_data_layout();
+        let dl = dl.borrow();
+        let mut align = if repr.pack.is_some() { dl.i8_align } else { dl.aggregate_align };
+        let mut max_repr_align = repr.align;
+
+        // If all the non-ZST fields have the same ABI and union ABI optimizations aren't
+        // disabled, we can use that common ABI for the union as a whole.
+        struct AbiMismatch;
+        let mut common_non_zst_abi_and_align = if repr.inhibit_union_abi_opt() {
+            // Can't optimize
+            Err(AbiMismatch)
+        } else {
+            Ok(None)
+        };
+
+        let mut size = Size::ZERO;
+        let only_variant = &variants[VariantIdx::new(0)];
+        for field in only_variant {
+            if field.is_unsized() {
+                self.delayed_bug("unsized field in union".to_string());
+            }
+
+            align = align.max(field.align);
+            max_repr_align = max_repr_align.max(field.max_repr_align);
+            size = cmp::max(size, field.size);
+
+            if field.is_zst() {
+                // Nothing more to do for ZST fields
+                continue;
+            }
+
+            if let Ok(common) = common_non_zst_abi_and_align {
+                // Discard valid range information and allow undef
+                let field_abi = field.abi.to_union();
+
+                if let Some((common_abi, common_align)) = common {
+                    if common_abi != field_abi {
+                        // Different fields have different ABI: disable opt
+                        common_non_zst_abi_and_align = Err(AbiMismatch);
+                    } else {
+                        // Fields with the same non-Aggregate ABI should also
+                        // have the same alignment
+                        if !matches!(common_abi, Abi::Aggregate { .. }) {
+                            assert_eq!(
+                                common_align, field.align.abi,
+                                "non-Aggregate field with matching ABI but differing alignment"
+                            );
+                        }
+                    }
+                } else {
+                    // First non-ZST field: record its ABI and alignment
+                    common_non_zst_abi_and_align = Ok(Some((field_abi, field.align.abi)));
+                }
+            }
+        }
+
+        if let Some(pack) = repr.pack {
+            align = align.min(AbiAndPrefAlign::new(pack));
+        }
+        // The unadjusted ABI alignment does not include repr(align), but does include repr(pack).
+        // See documentation on `LayoutS::unadjusted_abi_align`.
+        let unadjusted_abi_align = align.abi;
+        if let Some(repr_align) = repr.align {
+            align = align.max(AbiAndPrefAlign::new(repr_align));
+        }
+        // `align` must not be modified after this, or `unadjusted_abi_align` could be inaccurate.
+        let align = align;
+
+        // If all non-ZST fields have the same ABI, we may forward that ABI
+        // for the union as a whole, unless otherwise inhibited.
+        let abi = match common_non_zst_abi_and_align {
+            Err(AbiMismatch) | Ok(None) => Abi::Aggregate { sized: true },
+            Ok(Some((abi, _))) => {
+                if abi.inherent_align(dl).map(|a| a.abi) != Some(align.abi) {
+                    // Mismatched alignment (e.g. union is #[repr(packed)]): disable opt
+                    Abi::Aggregate { sized: true }
+                } else {
+                    abi
+                }
+            }
+        };
+
+        Some(LayoutS {
+            variants: Variants::Single { index: VariantIdx::new(0) },
+            fields: FieldsShape::Union(NonZeroUsize::new(only_variant.len())?),
+            abi,
+            largest_niche: None,
+            align,
+            size: size.align_to(align.abi),
+            max_repr_align,
+            unadjusted_abi_align,
+        })
+    }
+}
+
+/// single-variant enums are just structs, if you think about it
+fn layout_of_struct<'a, LC, FieldIdx: Idx, VariantIdx: Idx, F>(
+    layout_calc: &LC,
+    repr: &ReprOptions,
+    variants: &IndexSlice<VariantIdx, IndexVec<FieldIdx, F>>,
+    is_enum: bool,
+    is_unsafe_cell: bool,
+    scalar_valid_range: (Bound<u128>, Bound<u128>),
+    always_sized: bool,
+    dl: &TargetDataLayout,
+    present_first: VariantIdx,
+) -> Option<LayoutS<FieldIdx, VariantIdx>>
+where
+    LC: LayoutCalculator + ?Sized,
+    F: Deref<Target = &'a LayoutS<FieldIdx, VariantIdx>> + fmt::Debug,
+{
     // Struct, or univariant enum equivalent to a struct.
     // (Typechecking will reject discriminant-sizing attrs.)
 
@@ -213,7 +361,7 @@ pub trait LayoutCalculator {
         StructKind::MaybeUnsized
     };
 
-            let mut st = self.univariant(dl, &variants[v], repr, kind)?;
+    let mut st = layout_calc.univariant(dl, &variants[v], repr, kind)?;
     st.variants = Variants::Single { index: v };
 
     if is_unsafe_cell {
@@ -284,14 +432,22 @@ pub trait LayoutCalculator {
         ),
     }
 
-            return Some(st);
+    Some(st)
 }
 
-        // At this point, we have handled all unions and
+fn layout_of_enum<'a, LC, FieldIdx: Idx, VariantIdx: Idx, F>(
-        // structs. (We have also handled univariant enums
+    layout_calc: &LC,
-        // that allow representation optimization.)
+    repr: &ReprOptions,
-        assert!(is_enum);
+    variants: &IndexSlice<VariantIdx, IndexVec<FieldIdx, F>>,
-
+    discr_range_of_repr: impl Fn(i128, i128) -> (Integer, bool),
+    discriminants: impl Iterator<Item = (VariantIdx, i128)>,
+    dont_niche_optimize_enum: bool,
+    dl: &TargetDataLayout,
+) -> Option<LayoutS<FieldIdx, VariantIdx>>
+where
+    LC: LayoutCalculator + ?Sized,
+    F: Deref<Target = &'a LayoutS<FieldIdx, VariantIdx>> + fmt::Debug,
+{
     // Until we've decided whether to use the tagged or
     // niche filling LayoutS, we don't want to intern the
     // variant layouts, so we can't store them in the
@@ -318,7 +474,7 @@ pub trait LayoutCalculator {
         let mut variant_layouts = variants
             .iter_enumerated()
             .map(|(j, v)| {
-                    let mut st = self.univariant(dl, v, repr, StructKind::AlwaysSized)?;
+                let mut st = layout_calc.univariant(dl, v, repr, StructKind::AlwaysSized)?;
                 st.variants = Variants::Single { index: j };
 
                 align = align.max(st.align);
@@ -512,7 +668,7 @@ pub trait LayoutCalculator {
     let mut layout_variants = variants
         .iter_enumerated()
         .map(|(i, field_layouts)| {
-                let mut st = self.univariant(
+            let mut st = layout_calc.univariant(
                 dl,
                 field_layouts,
                 repr,
@@ -672,12 +828,14 @@ pub trait LayoutCalculator {
         }
         if let Some((prim, offset)) = common_prim {
             let prim_scalar = if common_prim_initialized_in_all_variants {
-                    scalar_unit(prim)
+                let size = prim.size(dl);
+                assert!(size.bits() <= 128);
+                Scalar::Initialized { value: prim, valid_range: WrappingRange::full(size) }
             } else {
                 // Common prim might be uninit.
                 Scalar::Union { value: prim }
             };
-                let pair = self.scalar_pair::<FieldIdx, VariantIdx>(tag, prim_scalar);
+            let pair = layout_calc.scalar_pair::<FieldIdx, VariantIdx>(tag, prim_scalar);
             let pair_offsets = match pair.fields {
                 FieldsShape::Arbitrary { ref offsets, ref memory_index } => {
                     assert_eq!(memory_index.raw, [0, 1]);
@@ -723,10 +881,7 @@ pub trait LayoutCalculator {
             tag_field: 0,
             variants: IndexVec::new(),
         },
-            fields: FieldsShape::Arbitrary {
+        fields: FieldsShape::Arbitrary { offsets: [Size::ZERO].into(), memory_index: [0].into() },
-                offsets: [Size::ZERO].into(),
-                memory_index: [0].into(),
-            },
         largest_niche,
         abi,
         align,
@@ -767,111 +922,6 @@ pub trait LayoutCalculator {
     Some(best_layout.layout)
 }
 
-    fn layout_of_union<
-        'a,
-        FieldIdx: Idx,
-        VariantIdx: Idx,
-        F: Deref<Target = &'a LayoutS<FieldIdx, VariantIdx>> + fmt::Debug,
-    >(
-        &self,
-        repr: &ReprOptions,
-        variants: &IndexSlice<VariantIdx, IndexVec<FieldIdx, F>>,
-    ) -> Option<LayoutS<FieldIdx, VariantIdx>> {
-        let dl = self.current_data_layout();
-        let dl = dl.borrow();
-        let mut align = if repr.pack.is_some() { dl.i8_align } else { dl.aggregate_align };
-        let mut max_repr_align = repr.align;
-
-        // If all the non-ZST fields have the same ABI and union ABI optimizations aren't
-        // disabled, we can use that common ABI for the union as a whole.
-        struct AbiMismatch;
-        let mut common_non_zst_abi_and_align = if repr.inhibit_union_abi_opt() {
-            // Can't optimize
-            Err(AbiMismatch)
-        } else {
-            Ok(None)
-        };
-
-        let mut size = Size::ZERO;
-        let only_variant = &variants[VariantIdx::new(0)];
-        for field in only_variant {
-            if field.is_unsized() {
-                self.delayed_bug("unsized field in union".to_string());
-            }
-
-            align = align.max(field.align);
-            max_repr_align = max_repr_align.max(field.max_repr_align);
-            size = cmp::max(size, field.size);
-
-            if field.is_zst() {
-                // Nothing more to do for ZST fields
-                continue;
-            }
-
-            if let Ok(common) = common_non_zst_abi_and_align {
-                // Discard valid range information and allow undef
-                let field_abi = field.abi.to_union();
-
-                if let Some((common_abi, common_align)) = common {
-                    if common_abi != field_abi {
-                        // Different fields have different ABI: disable opt
-                        common_non_zst_abi_and_align = Err(AbiMismatch);
-                    } else {
-                        // Fields with the same non-Aggregate ABI should also
-                        // have the same alignment
-                        if !matches!(common_abi, Abi::Aggregate { .. }) {
-                            assert_eq!(
-                                common_align, field.align.abi,
-                                "non-Aggregate field with matching ABI but differing alignment"
-                            );
-                        }
-                    }
-                } else {
-                    // First non-ZST field: record its ABI and alignment
-                    common_non_zst_abi_and_align = Ok(Some((field_abi, field.align.abi)));
-                }
-            }
-        }
-
-        if let Some(pack) = repr.pack {
-            align = align.min(AbiAndPrefAlign::new(pack));
-        }
-        // The unadjusted ABI alignment does not include repr(align), but does include repr(pack).
-        // See documentation on `LayoutS::unadjusted_abi_align`.
-        let unadjusted_abi_align = align.abi;
-        if let Some(repr_align) = repr.align {
-            align = align.max(AbiAndPrefAlign::new(repr_align));
-        }
-        // `align` must not be modified after this, or `unadjusted_abi_align` could be inaccurate.
-        let align = align;
-
-        // If all non-ZST fields have the same ABI, we may forward that ABI
-        // for the union as a whole, unless otherwise inhibited.
-        let abi = match common_non_zst_abi_and_align {
-            Err(AbiMismatch) | Ok(None) => Abi::Aggregate { sized: true },
-            Ok(Some((abi, _))) => {
-                if abi.inherent_align(dl).map(|a| a.abi) != Some(align.abi) {
-                    // Mismatched alignment (e.g. union is #[repr(packed)]): disable opt
-                    Abi::Aggregate { sized: true }
-                } else {
-                    abi
-                }
-            }
-        };
-
-        Some(LayoutS {
-            variants: Variants::Single { index: VariantIdx::new(0) },
-            fields: FieldsShape::Union(NonZeroUsize::new(only_variant.len())?),
-            abi,
-            largest_niche: None,
-            align,
-            size: size.align_to(align.abi),
-            max_repr_align,
-            unadjusted_abi_align,
-        })
-    }
-}
-
 /// Determines towards which end of a struct layout optimizations will try to place the best niches.
 enum NicheBias {
     Start,