Avoid alignment mismatch between ABI and layout for unions.

compiler/rustc_abi/src/layout.rs

@@ -729,42 +729,73 @@ pub trait LayoutCalculator {
  align = align.max(AbiAndPrefAlign::new(repr_align));
  }
 
- let optimize = !repr.inhibit_union_abi_opt();
+ // 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 mut abi = Abi::Aggregate { sized: true };
  let only_variant = &variants[FIRST_VARIANT];
  for field in only_variant {
  assert!(field.0.is_sized());
+
  align = align.max(field.align());
+ size = cmp::max(size, field.size());
 
- // If all non-ZST fields have the same ABI, forward this ABI
- if optimize && !field.0.is_zst() {
+ if field.0.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 = match field.abi() {
- Abi::Scalar(x) => Abi::Scalar(x.to_union()),
- Abi::ScalarPair(x, y) => Abi::ScalarPair(x.to_union(), y.to_union()),
- Abi::Vector { element: x, count } => {
- Abi::Vector { element: x.to_union(), count }
- }
- Abi::Uninhabited | Abi::Aggregate { .. } => Abi::Aggregate { sized: true },
- };
+ let field_abi = field.abi().to_union();
 
- if size == Size::ZERO {
- // first non ZST: initialize 'abi'
- abi = field_abi;
- } else if abi != field_abi {
- // different fields have different ABI: reset to Aggregate
- abi = Abi::Aggregate { sized: true };
+ 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)));
  }
  }
-
- size = cmp::max(size, field.size());
  }
 
  if let Some(pack) = repr.pack {
  align = align.min(AbiAndPrefAlign::new(pack));
  }
 
+ // 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: FIRST_VARIANT },
  fields: FieldsShape::Union(NonZeroUsize::new(only_variant.len())?),

compiler/rustc_abi/src/lib.rs

@@ -1272,6 +1272,50 @@ impl Abi {
  pub fn is_scalar(&self) -> bool {
  matches!(*self, Abi::Scalar(_))
  }
+
+ /// Returns the fixed alignment of this ABI, if any is mandated.
+ pub fn inherent_align<C: HasDataLayout>(&self, cx: &C) -> Option<AbiAndPrefAlign> {
+ Some(match *self {
+ Abi::Scalar(s) => s.align(cx),
+ Abi::ScalarPair(s1, s2) => s1.align(cx).max(s2.align(cx)),
+ Abi::Vector { element, count } => {
+ cx.data_layout().vector_align(element.size(cx) * count)
+ }
+ Abi::Uninhabited | Abi::Aggregate { .. } => return None,
+ })
+ }
+
+ /// Returns the fixed size of this ABI, if any is mandated.
+ pub fn inherent_size<C: HasDataLayout>(&self, cx: &C) -> Option<Size> {
+ Some(match *self {
+ Abi::Scalar(s) => {
+ // No padding in scalars.
+ s.size(cx)
+ }
+ Abi::ScalarPair(s1, s2) => {
+ // May have some padding between the pair.
+ let field2_offset = s1.size(cx).align_to(s2.align(cx).abi);
+ (field2_offset + s2.size(cx)).align_to(self.inherent_align(cx)?.abi)
+ }
+ Abi::Vector { element, count } => {
+ // No padding in vectors, except possibly for trailing padding
+ // to make the size a multiple of align (e.g. for vectors of size 3).
+ (element.size(cx) * count).align_to(self.inherent_align(cx)?.abi)
+ }
+ Abi::Uninhabited | Abi::Aggregate { .. } => return None,
+ })
+ }
+
+ /// Discard validity range information and allow undef.
+ pub fn to_union(&self) -> Self {
+ assert!(self.is_sized());
+ match *self {
+ Abi::Scalar(s) => Abi::Scalar(s.to_union()),
+ Abi::ScalarPair(s1, s2) => Abi::ScalarPair(s1.to_union(), s2.to_union()),
+ Abi::Vector { element, count } => Abi::Vector { element: element.to_union(), count },
+ Abi::Uninhabited | Abi::Aggregate { .. } => Abi::Aggregate { sized: true },
+ }
+ }
 }
 
 #[derive(PartialEq, Eq, Hash, Clone, Debug)]

compiler/rustc_ty_utils/src/layout_sanity_check.rs

@@ -4,7 +4,7 @@ use rustc_middle::ty::{
 };
 use rustc_target::abi::*;
 
-use std::cmp;
+use std::assert_matches::assert_matches;
 
 /// Enforce some basic invariants on layouts.
 pub(super) fn sanity_check_layout<'tcx>(
@@ -68,21 +68,31 @@ pub(super) fn sanity_check_layout<'tcx>(
  }
 
  fn check_layout_abi<'tcx>(cx: &LayoutCx<'tcx, TyCtxt<'tcx>>, layout: &TyAndLayout<'tcx>) {
+ // Verify the ABI mandated alignment and size.
+ let align = layout.abi.inherent_align(cx).map(|align| align.abi);
+ let size = layout.abi.inherent_size(cx);
+ let Some((align, size)) = align.zip(size) else {
+ assert_matches!(
+ layout.layout.abi(),
+ Abi::Uninhabited | Abi::Aggregate { .. },
+ "ABI unexpectedly missing alignment and/or size in {layout:#?}"
+ );
+ return
+ };
+ assert_eq!(
+ layout.layout.align().abi,
+ align,
+ "alignment mismatch between ABI and layout in {layout:#?}"
+ );
+ assert_eq!(
+ layout.layout.size(),
+ size,
+ "size mismatch between ABI and layout in {layout:#?}"
+ );
+
+ // Verify per-ABI invariants
  match layout.layout.abi() {
- Abi::Scalar(scalar) => {
- // No padding in scalars.
- let size = scalar.size(cx);
- let align = scalar.align(cx).abi;
- assert_eq!(
- layout.layout.size(),
- size,
- "size mismatch between ABI and layout in {layout:#?}"
- );
- assert_eq!(
- layout.layout.align().abi,
- align,
- "alignment mismatch between ABI and layout in {layout:#?}"
- );
+ Abi::Scalar(_) => {
  // Check that this matches the underlying field.
  let inner = skip_newtypes(cx, layout);
  assert!(
@@ -135,24 +145,6 @@ pub(super) fn sanity_check_layout<'tcx>(
  }
  }
  Abi::ScalarPair(scalar1, scalar2) => {
- // Sanity-check scalar pairs. Computing the expected size and alignment is a bit of work.
- let size1 = scalar1.size(cx);
- let align1 = scalar1.align(cx).abi;
- let size2 = scalar2.size(cx);
- let align2 = scalar2.align(cx).abi;
- let align = cmp::max(align1, align2);
- let field2_offset = size1.align_to(align2);
- let size = (field2_offset + size2).align_to(align);
- assert_eq!(
- layout.layout.size(),
- size,
- "size mismatch between ABI and layout in {layout:#?}"
- );
- assert_eq!(
- layout.layout.align().abi,
- align,
- "alignment mismatch between ABI and layout in {layout:#?}",
- );
  // Check that the underlying pair of fields matches.
  let inner = skip_newtypes(cx, layout);
  assert!(
@@ -189,8 +181,9 @@ pub(super) fn sanity_check_layout<'tcx>(
  "`ScalarPair` layout for type with less than two non-ZST fields: {inner:#?}"
  )
  });
- assert!(
- fields.next().is_none(),
+ assert_matches!(
+ fields.next(),
+ None,
  "`ScalarPair` layout for type with at least three non-ZST fields: {inner:#?}"
  );
  // The fields might be in opposite order.
@@ -200,6 +193,10 @@ pub(super) fn sanity_check_layout<'tcx>(
  (offset2, field2, offset1, field1)
  };
  // The fields should be at the right offset, and match the `scalar` layout.
+ let size1 = scalar1.size(cx);
+ let align1 = scalar1.align(cx).abi;
+ let size2 = scalar2.size(cx);
+ let align2 = scalar2.align(cx).abi;
  assert_eq!(
  offset1,
  Size::ZERO,
@@ -213,10 +210,12 @@ pub(super) fn sanity_check_layout<'tcx>(
  field1.align.abi, align1,
  "`ScalarPair` first field with bad align in {inner:#?}",
  );
- assert!(
- matches!(field1.abi, Abi::Scalar(_)),
+ assert_matches!(
+ field1.abi,
+ Abi::Scalar(_),
  "`ScalarPair` first field with bad ABI in {inner:#?}",
  );
+ let field2_offset = size1.align_to(align2);
  assert_eq!(
  offset2, field2_offset,
  "`ScalarPair` second field at bad offset in {inner:#?}",
@@ -229,27 +228,14 @@ pub(super) fn sanity_check_layout<'tcx>(
  field2.align.abi, align2,
  "`ScalarPair` second field with bad align in {inner:#?}",
  );
- assert!(
- matches!(field2.abi, Abi::Scalar(_)),
+ assert_matches!(
+ field2.abi,
+ Abi::Scalar(_),
  "`ScalarPair` second field with bad ABI in {inner:#?}",
  );
  }
- Abi::Vector { count, element } => {
- // No padding in vectors, except possibly for trailing padding to make the size a multiple of align.
- let size = element.size(cx) * count;
- let align = cx.data_layout().vector_align(size).abi;
- let size = size.align_to(align); // needed e.g. for vectors of size 3
+ Abi::Vector { element, .. } => {
  assert!(align >= element.align(cx).abi); // just sanity-checking `vector_align`.
- assert_eq!(
- layout.layout.size(),
- size,
- "size mismatch between ABI and layout in {layout:#?}"
- );
- assert_eq!(
- layout.layout.align().abi,
- align,
- "alignment mismatch between ABI and layout in {layout:#?}"
- );
  // FIXME: Do some kind of check of the inner type, like for Scalar and ScalarPair.
  }
  Abi::Uninhabited | Abi::Aggregate { .. } => {} // Nothing to check.

compiler/rustc_ty_utils/src/lib.rs

@@ -5,6 +5,7 @@
 //! This API is completely unstable and subject to change.
 
 #![doc(html_root_url = "https://doc.rust-lang.org/nightly/nightly-rustc/")]
+#![feature(assert_matches)]
 #![feature(iterator_try_collect)]
 #![feature(let_chains)]
 #![feature(never_type)]

tests/ui/layout/debug.rs

@@ -1,8 +1,9 @@
 // normalize-stderr-test "pref: Align\([1-8] bytes\)" -> "pref: $$PREF_ALIGN"
-#![feature(never_type, rustc_attrs, type_alias_impl_trait)]
+#![feature(never_type, rustc_attrs, type_alias_impl_trait, repr_simd)]
 #![crate_type = "lib"]
 
 #[rustc_layout(debug)]
+#[derive(Copy, Clone)]
 enum E { Foo, Bar(!, i32, i32) } //~ ERROR: layout_of
 
 #[rustc_layout(debug)]
@@ -17,6 +18,51 @@ type Test = Result<i32, i32>; //~ ERROR: layout_of
 #[rustc_layout(debug)]
 type T = impl std::fmt::Debug; //~ ERROR: layout_of
 
+#[rustc_layout(debug)]
+pub union V { //~ ERROR: layout_of
+ a: [u16; 0],
+ b: u8,
+}
+
+#[rustc_layout(debug)]
+pub union W { //~ ERROR: layout_of
+ b: u8,
+ a: [u16; 0],
+}
+
+#[rustc_layout(debug)]
+pub union Y { //~ ERROR: layout_of
+ b: [u8; 0],
+ a: [u16; 0],
+}
+
+#[rustc_layout(debug)]
+#[repr(packed(1))]
+union P1 { x: u32 } //~ ERROR: layout_of
+
+#[rustc_layout(debug)]
+#[repr(packed(1))]
+union P2 { x: (u32, u32) } //~ ERROR: layout_of
+
+#[repr(simd)]
+#[derive(Copy, Clone)]
+struct F32x4(f32, f32, f32, f32);
+
+#[rustc_layout(debug)]
+#[repr(packed(1))]
+union P3 { x: F32x4 } //~ ERROR: layout_of
+
+#[rustc_layout(debug)]
+#[repr(packed(1))]
+union P4 { x: E } //~ ERROR: layout_of
+
+#[rustc_layout(debug)]
+#[repr(packed(1))]
+union P5 { zst: [u16; 0], byte: u8 } //~ ERROR: layout_of
+
+#[rustc_layout(debug)]
+type X = std::mem::MaybeUninit<u8>; //~ ERROR: layout_of
+
 fn f() -> T {
  0i32
 }