Adjust the alignment when passing the niche as a pointer

Author: dianqk

compiler/rustc_middle/src/ty/layout.rs

@@ -999,7 +999,7 @@ where
             }
 
             _ => {
-                let mut data_variant = match this.variants {
+                let (mut data_variant, 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.
@@ -1012,13 +1012,32 @@ where
                     // using more niches than just null (e.g., the first page of
                     // the address space, or unaligned pointers).
                     Variants::Multiple {
-                        tag_encoding: TagEncoding::Niche { untagged_variant, .. },
+                        tag_encoding:
+                            TagEncoding::Niche { untagged_variant, ref niche_variants, niche_start },
                         tag_field,
+                        ref variants,
                         ..
                     } 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,
+                        // so we need to adjust the alignment to an appropriate value.
+                        let align = variants
+                            .iter_enumerated()
+                            .filter_map(|(idx, variant)| {
+                                if untagged_variant == idx || variant.abi.is_uninhabited() {
+                                    None
+                                } else {
+                                    Some(idx)
+                                }
+                            })
+                            .fold(Align::MAX, |align, idx| {
+                                let niche = (idx.as_u32() as u128
+                                    - niche_variants.start().as_u32() as u128)
+                                    .wrapping_add(niche_start);
+                                align.restrict_for_offset(Size::from_bytes(niche))
+                            });
+                        (Some(this.for_variant(cx, untagged_variant)), Some(align))
                     }
-                    _ => Some(this),
+                    _ => (Some(this), None),
                 };
 
                 if let Some(variant) = data_variant {
@@ -1055,20 +1074,23 @@ 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) = align {
+                        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) {}