Adjust the alignment when passing a niche as a pointer

Author: dianqk

compiler/rustc_abi/src/lib.rs

@@ -1489,6 +1489,29 @@ pub enum Variants<FieldIdx: Idx, VariantIdx: Idx> {
     },
 }
 
+impl<FieldIdx: Idx, VariantIdx: Idx> Variants<FieldIdx, VariantIdx> {
+    // Returns niches for the discriminants.
+    pub fn niches(&self) -> Option<impl Iterator<Item = u128> + '_> {
+        match self {
+            &Variants::Multiple {
+                tag_encoding:
+                    TagEncoding::Niche { untagged_variant, ref niche_variants, niche_start },
+                ref variants,
+                ..
+            } => Some(variants.iter_enumerated().filter_map(move |(variant_idx, variant)| {
+                if untagged_variant == variant_idx || variant.abi.is_uninhabited() {
+                    None
+                } else {
+                    let niche = ((variant_idx.index() - niche_variants.start().index()) as u128)
+                        .wrapping_add(niche_start);
+                    Some(niche)
+                }
+            })),
+            _ => None,
+        }
+    }
+}
+
 // NOTE: This struct is generic over the VariantIdx for rust-analyzer usage.
 #[derive(PartialEq, Eq, Hash, Clone, Debug)]
 #[cfg_attr(feature = "nightly", derive(HashStable_Generic))]

compiler/rustc_middle/src/ty/layout.rs

@@ -999,7 +999,7 @@ where
             }
 
             _ => {
-                let mut data_variant = match this.variants {
+                let (mut data_variant, niche_align) = match this.variants {
                     // Within the discriminant field, only the niche itself is
                     // always initialized, so we only check for a pointer at its
                     // offset.
@@ -1016,9 +1016,16 @@ where
                         tag_field,
                         ..
                     } if this.fields.offset(tag_field) == offset => {
-                        Some(this.for_variant(cx, untagged_variant))
+                        // When a non-null niche is passed in, we might pass an unaligned value.
+                        // Calculates a maximum alignment that matches all niches.
+                        let niche_align = this.variants.niches().map(|niches| {
+                            niches.fold(Align::MAX, |align, niche| {
+                                align.restrict_for_offset(Size::from_bytes(niche))
+                            })
+                        });
+                        (Some(this.for_variant(cx, untagged_variant)), niche_align)
                     }
-                    _ => Some(this),
+                    _ => (Some(this), None),
                 };
 
                 if let Some(variant) = data_variant {
@@ -1055,20 +1062,24 @@ where
                     }
                 }
 
-                // Fixup info for the first field of a `Box`. Recursive traversal will have found
-                // the raw pointer, so size and align are set to the boxed type, but `pointee.safe`
-                // will still be `None`.
                 if let Some(ref mut pointee) = result {
                     if offset.bytes() == 0
                         && let Some(boxed_ty) = this.ty.boxed_ty()
                     {
+                        // Fixup info for the first field of a `Box`. Recursive traversal will have found
+                        // the raw pointer, so size and align are set to the boxed type, but `pointee.safe`
+                        // will still be `None`.
                         debug_assert!(pointee.safe.is_none());
                         let optimize = tcx.sess.opts.optimize != OptLevel::No;
                         pointee.safe = Some(PointerKind::Box {
                             unpin: optimize && boxed_ty.is_unpin(tcx, cx.param_env()),
                             global: this.ty.is_box_global(tcx),
                         });
                     }
+                    if let Some(align) = niche_align {
+                        // Takes the minimum alignment of the pointer and niches to ensure that the alignment is legal.
+                        pointee.align = pointee.align.min(align);
+                    }
                 }
 
                 result

tests/codegen/enum/enum-ptr-align-niche.rs

@@ -0,0 +1,113 @@
+//@ compile-flags: -Cno-prepopulate-passes -O
+//@ only-64bit (I don't care about alignment under different bits)
+
+// Testing different niches updates to the corresponding alignment.
+
+#![crate_type = "lib"]
+#![feature(rustc_attrs)]
+#![feature(never_type)]
+
+#[rustc_layout_scalar_valid_range_start(0)]
+#[rustc_layout_scalar_valid_range_end(0x7fff)]
+struct RestrictedAddress_0_0x7fff(&'static i64);
+
+#[rustc_layout_scalar_valid_range_start(1)]
+#[rustc_layout_scalar_valid_range_end(0x7fff)]
+struct RestrictedAddress_1_0x7fff(&'static i64);
+
+#[rustc_layout_scalar_valid_range_start(0)]
+#[rustc_layout_scalar_valid_range_end(0xf000)]
+struct RestrictedAddress_0_0xf000(&'static i64);
+
+enum MultipleAlign8 {
+    Untag(RestrictedAddress_1_0x7fff),
+    Niche_32768,
+    Uninhabited_1(!),
+    Uninhabited_2(!),
+    Uninhabited_3(!),
+    Uninhabited_4(!),
+    Uninhabited_5(!),
+    Uninhabited_6(!),
+    Uninhabited_7(!),
+    Niche_32776,
+}
+
+// CHECK-LABEL: @multiple_niches_align_8
+// CHECK-SAME: align 8 {{.*}}%a)
+#[no_mangle]
+#[inline(never)]
+fn multiple_niches_align_8(a: MultipleAlign8) {}
+
+// CHECK-LABEL: @call_multiple_niches_align_8
+#[no_mangle]
+fn call_multiple_niches_align_8() {
+    // CHECK: call void @multiple_niches_align_8(ptr {{.*}}align 8 {{.*}}(i64 32768 to ptr)
+    multiple_niches_align_8(MultipleAlign8::Niche_32768);
+    // CHECK: call void @multiple_niches_align_8(ptr {{.*}}align 8 {{.*}}(i64 32776 to ptr)
+    multiple_niches_align_8(MultipleAlign8::Niche_32776);
+}
+
+enum MultipleAlign2 {
+    Untag(RestrictedAddress_1_0x7fff),
+    Niche_32768,
+    Uninhabited_1(!),
+    Niche_32770,
+}
+
+// CHECK-LABEL: @multiple_niches_align_2
+// CHECK-SAME: align 2 {{.*}}%a)
+#[no_mangle]
+#[inline(never)]
+fn multiple_niches_align_2(a: MultipleAlign2) {}
+
+// CHECK-LABEL: @call_multiple_niches_align_2
+#[no_mangle]
+fn call_multiple_niches_align_2() {
+    // CHECK: call void @multiple_niches_align_2(ptr {{.*}}align 2 {{.*}}(i64 32768 to ptr)
+    multiple_niches_align_2(MultipleAlign2::Niche_32768);
+    // CHECK: call void @multiple_niches_align_2(ptr {{.*}}align 2 {{.*}}(i64 32770 to ptr)
+    multiple_niches_align_2(MultipleAlign2::Niche_32770);
+}
+
+enum SingleAlign8 {
+    Untag(RestrictedAddress_0_0x7fff),
+    Niche_32768,
+}
+
+// CHECK-LABEL: @single_niche_align_8
+// CHECK-SAME: align 8 {{.*}}%a)
+#[no_mangle]
+#[inline(never)]
+fn single_niche_align_8(a: SingleAlign8) {}
+
+// CHECK-LABEL: @call_single_niche_align_8
+#[no_mangle]
+fn call_single_niche_align_8() {
+    // CHECK: call void @single_niche_align_8(ptr {{.*}}align 8 {{.*}}(i64 32768 to ptr)
+    single_niche_align_8(SingleAlign8::Niche_32768);
+}
+
+enum SingleAlign1 {
+    Untag(RestrictedAddress_0_0xf000),
+    Niche_61441,
+}
+
+// CHECK-LABEL: @single_niche_align_1
+// CHECK-SAME: align 1 {{.*}}%a)
+#[no_mangle]
+#[inline(never)]
+fn single_niche_align_1(a: SingleAlign1) {}
+
+// CHECK-LABEL: @call_single_niche_align_1
+#[no_mangle]
+fn call_single_niche_align_1() {
+    // CHECK: call void @single_niche_align_1(ptr {{.*}}align 1 {{.*}}(i64 61441 to ptr)
+    single_niche_align_1(SingleAlign1::Niche_61441);
+}
+
+// Check that we only apply the new alignment on enum.
+
+// CHECK-LABEL: @restricted_address
+// CHECK-SAME: align 8 {{.*}}%a)
+#[no_mangle]
+fn restricted_address(a: RestrictedAddress_0_0x7fff) {}