Move more layouting logic to rustc_abi

compiler/rustc_abi/src/layout.rs

@@ -4,6 +4,7 @@ use std::{cmp, iter};
 
 use rustc_hashes::Hash64;
 use rustc_index::Idx;
+use rustc_index::bit_set::BitMatrix;
 use tracing::debug;
 
 use crate::{
@@ -12,6 +13,9 @@ use crate::{
     Variants, WrappingRange,
 };
 
+mod coroutine;
+mod simple;
+
 #[cfg(feature = "nightly")]
 mod ty;
 
@@ -60,31 +64,44 @@ pub enum LayoutCalculatorError<F> {
 
     /// The fields or variants have irreconcilable reprs
     ReprConflict,
+
+    /// The length of an SIMD type is zero
+    ZeroLengthSimdType,
+
+    /// The length of an SIMD type exceeds the maximum number of lanes
+    OversizedSimdType { max_lanes: u64 },
+
+    /// An element type of an SIMD type isn't a primitive
+    NonPrimitiveSimdType(F),
 }
 
 impl<F> LayoutCalculatorError<F> {
     pub fn without_payload(&self) -> LayoutCalculatorError<()> {
-        match self {
-            LayoutCalculatorError::UnexpectedUnsized(_) => {
-                LayoutCalculatorError::UnexpectedUnsized(())
-            }
-            LayoutCalculatorError::SizeOverflow => LayoutCalculatorError::SizeOverflow,
-            LayoutCalculatorError::EmptyUnion => LayoutCalculatorError::EmptyUnion,
-            LayoutCalculatorError::ReprConflict => LayoutCalculatorError::ReprConflict,
+        use LayoutCalculatorError::*;
+        match *self {
+            UnexpectedUnsized(_) => UnexpectedUnsized(()),
+            SizeOverflow => SizeOverflow,
+            EmptyUnion => EmptyUnion,
+            ReprConflict => ReprConflict,
+            ZeroLengthSimdType => ZeroLengthSimdType,
+            OversizedSimdType { max_lanes } => OversizedSimdType { max_lanes },
+            NonPrimitiveSimdType(_) => NonPrimitiveSimdType(()),
         }
     }
 
     /// Format an untranslated diagnostic for this type
     ///
     /// Intended for use by rust-analyzer, as neither it nor `rustc_abi` depend on fluent infra.
     pub fn fallback_fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        use LayoutCalculatorError::*;
         f.write_str(match self {
-            LayoutCalculatorError::UnexpectedUnsized(_) => {
-                "an unsized type was found where a sized type was expected"
+            UnexpectedUnsized(_) => "an unsized type was found where a sized type was expected",
+            SizeOverflow => "size overflow",
+            EmptyUnion => "type is a union with no fields",
+            ReprConflict => "type has an invalid repr",
+            ZeroLengthSimdType | OversizedSimdType { .. } | NonPrimitiveSimdType(_) => {
+                "invalid simd type definition"
             }
-            LayoutCalculatorError::SizeOverflow => "size overflow",
-            LayoutCalculatorError::EmptyUnion => "type is a union with no fields",
-            LayoutCalculatorError::ReprConflict => "type has an invalid repr",
         })
     }
 }
@@ -102,41 +119,115 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
         Self { cx }
     }
 
-    pub fn scalar_pair<FieldIdx: Idx, VariantIdx: Idx>(
+    pub fn array_like<FieldIdx: Idx, VariantIdx: Idx, F>(
         &self,
-        a: Scalar,
-        b: Scalar,
-    ) -> LayoutData<FieldIdx, VariantIdx> {
-        let dl = self.cx.data_layout();
-        let b_align = b.align(dl);
-        let align = a.align(dl).max(b_align).max(dl.aggregate_align);
-        let b_offset = a.size(dl).align_to(b_align.abi);
-        let size = (b_offset + b.size(dl)).align_to(align.abi);
+        element: &LayoutData<FieldIdx, VariantIdx>,
+        count_if_sized: Option<u64>, // None for slices
+    ) -> LayoutCalculatorResult<FieldIdx, VariantIdx, F> {
+        let count = count_if_sized.unwrap_or(0);
+        let size =
+            element.size.checked_mul(count, &self.cx).ok_or(LayoutCalculatorError::SizeOverflow)?;
 
-        // HACK(nox): We iter on `b` and then `a` because `max_by_key`
-        // returns the last maximum.
-        let largest_niche = Niche::from_scalar(dl, b_offset, b)
-            .into_iter()
-            .chain(Niche::from_scalar(dl, Size::ZERO, a))
-            .max_by_key(|niche| niche.available(dl));
+        Ok(LayoutData {
+            variants: Variants::Single { index: VariantIdx::new(0) },
+            fields: FieldsShape::Array { stride: element.size, count },
+            backend_repr: BackendRepr::Memory { sized: count_if_sized.is_some() },
+            largest_niche: element.largest_niche.filter(|_| count != 0),
+            uninhabited: element.uninhabited && count != 0,
+            align: element.align,
+            size,
+            max_repr_align: None,
+            unadjusted_abi_align: element.align.abi,
+            randomization_seed: element.randomization_seed.wrapping_add(Hash64::new(count)),
+        })
+    }
 
-        let combined_seed = a.size(&self.cx).bytes().wrapping_add(b.size(&self.cx).bytes());
+    pub fn simd_type<
+        FieldIdx: Idx,
+        VariantIdx: Idx,
+        F: AsRef<LayoutData<FieldIdx, VariantIdx>> + fmt::Debug,
+    >(
+        &self,
+        element: F,
+        count: u64,
+        repr_packed: bool,
+    ) -> LayoutCalculatorResult<FieldIdx, VariantIdx, F> {
+        let elt = element.as_ref();
+        if count == 0 {
+            return Err(LayoutCalculatorError::ZeroLengthSimdType);
+        } else if count > crate::MAX_SIMD_LANES {
+            return Err(LayoutCalculatorError::OversizedSimdType {
+                max_lanes: crate::MAX_SIMD_LANES,
+            });
+        }
 
-        LayoutData {
+        let BackendRepr::Scalar(e_repr) = elt.backend_repr else {
+            return Err(LayoutCalculatorError::NonPrimitiveSimdType(element));
+        };
+
+        // Compute the size and alignment of the vector
+        let dl = self.cx.data_layout();
+        let size =
+            elt.size.checked_mul(count, dl).ok_or_else(|| LayoutCalculatorError::SizeOverflow)?;
+        let (repr, align) = if repr_packed && !count.is_power_of_two() {
+            // Non-power-of-two vectors have padding up to the next power-of-two.
+            // If we're a packed repr, remove the padding while keeping the alignment as close
+            // to a vector as possible.
+            (
+                BackendRepr::Memory { sized: true },
+                AbiAndPrefAlign {
+                    abi: Align::max_aligned_factor(size),
+                    pref: dl.llvmlike_vector_align(size).pref,
+                },
+            )
+        } else {
+            (BackendRepr::SimdVector { element: e_repr, count }, dl.llvmlike_vector_align(size))
+        };
+        let size = size.align_to(align.abi);
+
+        Ok(LayoutData {
             variants: Variants::Single { index: VariantIdx::new(0) },
             fields: FieldsShape::Arbitrary {
-                offsets: [Size::ZERO, b_offset].into(),
-                memory_index: [0, 1].into(),
+                offsets: [Size::ZERO].into(),
+                memory_index: [0].into(),
             },
-            backend_repr: BackendRepr::ScalarPair(a, b),
-            largest_niche,
+            backend_repr: repr,
+            largest_niche: elt.largest_niche,
             uninhabited: false,
-            align,
             size,
+            align,
             max_repr_align: None,
-            unadjusted_abi_align: align.abi,
-            randomization_seed: Hash64::new(combined_seed),
-        }
+            unadjusted_abi_align: elt.align.abi,
+            randomization_seed: elt.randomization_seed.wrapping_add(Hash64::new(count)),
+        })
+    }
+
+    /// Compute the layout for a coroutine.
+    ///
+    /// This uses dedicated code instead of [`Self::layout_of_struct_or_enum`], as coroutine
+    /// fields may be shared between multiple variants (see the [`coroutine`] module for details).
+    pub fn coroutine<
+        'a,
+        F: Deref<Target = &'a LayoutData<FieldIdx, VariantIdx>> + fmt::Debug + Copy,
+        VariantIdx: Idx,
+        FieldIdx: Idx,
+        LocalIdx: Idx,
+    >(
+        &self,
+        local_layouts: &IndexSlice<LocalIdx, F>,
+        prefix_layouts: IndexVec<FieldIdx, F>,
+        variant_fields: &IndexSlice<VariantIdx, IndexVec<FieldIdx, LocalIdx>>,
+        storage_conflicts: &BitMatrix<LocalIdx, LocalIdx>,
+        tag_to_layout: impl Fn(Scalar) -> F,
+    ) -> LayoutCalculatorResult<FieldIdx, VariantIdx, F> {
+        coroutine::layout(
+            self,
+            local_layouts,
+            prefix_layouts,
+            variant_fields,
+            storage_conflicts,
+            tag_to_layout,
+        )
     }
 
     pub fn univariant<
@@ -214,25 +305,6 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
         layout
     }
 
-    pub fn layout_of_never_type<FieldIdx: Idx, VariantIdx: Idx>(
-        &self,
-    ) -> LayoutData<FieldIdx, VariantIdx> {
-        let dl = self.cx.data_layout();
-        // This is also used for uninhabited enums, so we use `Variants::Empty`.
-        LayoutData {
-            variants: Variants::Empty,
-            fields: FieldsShape::Primitive,
-            backend_repr: BackendRepr::Memory { sized: true },
-            largest_niche: None,
-            uninhabited: true,
-            align: dl.i8_align,
-            size: Size::ZERO,
-            max_repr_align: None,
-            unadjusted_abi_align: dl.i8_align.abi,
-            randomization_seed: Hash64::ZERO,
-        }
-    }
-
     pub fn layout_of_struct_or_enum<
         'a,
         FieldIdx: Idx,
@@ -260,7 +332,7 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
             Some(present_first) => present_first,
             // Uninhabited because it has no variants, or only absent ones.
             None if is_enum => {
-                return Ok(self.layout_of_never_type());
+                return Ok(LayoutData::never_type(&self.cx));
             }
             // If it's a struct, still compute a layout so that we can still compute the
             // field offsets.
@@ -949,7 +1021,8 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                     // Common prim might be uninit.
                     Scalar::Union { value: prim }
                 };
-                let pair = self.scalar_pair::<FieldIdx, VariantIdx>(tag, prim_scalar);
+                let pair =
+                    LayoutData::<FieldIdx, VariantIdx>::scalar_pair(&self.cx, tag, prim_scalar);
                 let pair_offsets = match pair.fields {
                     FieldsShape::Arbitrary { ref offsets, ref memory_index } => {
                         assert_eq!(memory_index.raw, [0, 1]);
@@ -1341,7 +1414,8 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                             } else {
                                 ((j, b), (i, a))
                             };
-                            let pair = self.scalar_pair::<FieldIdx, VariantIdx>(a, b);
+                            let pair =
+                                LayoutData::<FieldIdx, VariantIdx>::scalar_pair(&self.cx, a, b);
                             let pair_offsets = match pair.fields {
                                 FieldsShape::Arbitrary { ref offsets, ref memory_index } => {
                                     assert_eq!(memory_index.raw, [0, 1]);