Fix the size of niche enums with ZST alignment

For enums with an aligned ZST variant, like `[T; 0]`, the niche layout
was not computing a sufficient size to be consistent with alignment. Now
we pad that size up to the alignment, and also make sure to only use the
niche variant's ABI when the size and alignment still match.

3 files changed, 348 insertions, 2 deletions

diff --git a/compiler/rustc_middle/src/ty/layout.rs b/compiler/rustc_middle/src/ty/layout.rs
index dde55dd9655..833edd22805 100644
--- a/compiler/rustc_middle/src/ty/layout.rs
+++ b/compiler/rustc_middle/src/ty/layout.rs
@@ -1231,11 +1231,15 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
                                 .collect::<Result<IndexVec<VariantIdx, _>, _>>()?;
 
                             let offset = st[i].fields().offset(field_index) + niche.offset;
-                            let size = st[i].size();
+
+                            // Align the total size to the largest alignment.
+                            let size = st[i].size().align_to(align.abi);
 
                             let abi = if st.iter().all(|v| v.abi().is_uninhabited()) {
                                 Abi::Uninhabited
-                            } else {
+                            } else if align == st[i].align() && size == st[i].size() {
+                                // When the total alignment and size match, we can use the
+                                // same ABI as the scalar variant with the reserved niche.
                                 match st[i].abi() {
                                     Abi::Scalar(_) => Abi::Scalar(niche_scalar),
                                     Abi::ScalarPair(first, second) => {
@@ -1249,6 +1253,8 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
                                     }
                                     _ => Abi::Aggregate { sized: true },
                                 }
+                            } else {
+                                Abi::Aggregate { sized: true }
                             };
 
                             let largest_niche = Niche::from_scalar(dl, offset, niche_scalar);
diff --git a/src/test/ui/layout/zero-sized-array-enum-niche.rs b/src/test/ui/layout/zero-sized-array-enum-niche.rs
new file mode 100644
index 00000000000..c03077faed5
--- /dev/null
+++ b/src/test/ui/layout/zero-sized-array-enum-niche.rs
@@ -0,0 +1,37 @@
+// normalize-stderr-test "pref: Align\([1-8] bytes\)" -> "pref: $$PREF_ALIGN"
+#![feature(rustc_attrs)]
+#![crate_type = "lib"]
+
+// Various tests around the behavior of zero-sized arrays and
+// enum niches, especially that they have coherent size and alignment.
+
+// The original problem in #99836 came from ndarray's `TryFrom` for
+// `SliceInfo<[SliceInfoElem; 0], Din, Dout>`, where that returns
+// `Result<Self, ShapeError>` ~= `Result<AlignedZST, TypeWithNiche>`.
+// This is a close enough approximation:
+#[rustc_layout(debug)]
+type AlignedResult = Result<[u32; 0], bool>; //~ ERROR: layout_of
+// The bug gave that size 1 with align 4, but the size should also be 4.
+// It was also using the bool niche for the enum tag, which is fine, but
+// after the fix, layout decides to use a direct tagged repr instead.
+
+
+// Tagged repr is clever enough to grow tags to fill any padding, e.g.:
+// 1.   `T_FF` (one byte of Tag, one byte of padding, two bytes of align=2 Field)
+//   -> `TTFF` (Tag has expanded to two bytes, i.e. like `#[repr(u16)]`)
+// 2.    `TFF` (one byte of Tag, two bytes of align=1 Field)
+//   -> Tag has no room to expand!
+//   (this outcome can be forced onto 1. by wrapping Field in `Packed<...>`)
+#[repr(packed)]
+struct Packed<T>(T);
+
+#[rustc_layout(debug)]
+type NicheLosesToTagged = Result<[u32; 0], Packed<std::num::NonZeroU16>>; //~ ERROR: layout_of
+// Should get tag_encoding: Direct, size == align == 4.
+
+#[repr(u16)]
+enum U16IsZero { _Zero = 0 }
+
+#[rustc_layout(debug)]
+type NicheWinsOverTagged = Result<[u32; 0], Packed<U16IsZero>>; //~ ERROR: layout_of
+// Should get tag_encoding: Niche, size == align == 4.
diff --git a/src/test/ui/layout/zero-sized-array-enum-niche.stderr b/src/test/ui/layout/zero-sized-array-enum-niche.stderr
new file mode 100644
index 00000000000..9fae1974db8
--- /dev/null
+++ b/src/test/ui/layout/zero-sized-array-enum-niche.stderr
@@ -0,0 +1,303 @@
+error: layout_of(std::result::Result<[u32; 0], bool>) = Layout {
+           fields: Arbitrary {
+               offsets: [
+                   Size(0 bytes),
+               ],
+               memory_index: [
+                   0,
+               ],
+           },
+           variants: Multiple {
+               tag: Initialized {
+                   value: Int(
+                       I8,
+                       false,
+                   ),
+                   valid_range: 0..=1,
+               },
+               tag_encoding: Direct,
+               tag_field: 0,
+               variants: [
+                   Layout {
+                       fields: Arbitrary {
+                           offsets: [
+                               Size(4 bytes),
+                           ],
+                           memory_index: [
+                               0,
+                           ],
+                       },
+                       variants: Single {
+                           index: 0,
+                       },
+                       abi: Aggregate {
+                           sized: true,
+                       },
+                       largest_niche: None,
+                       align: AbiAndPrefAlign {
+                           abi: Align(4 bytes),
+                           pref: $PREF_ALIGN,
+                       },
+                       size: Size(4 bytes),
+                   },
+                   Layout {
+                       fields: Arbitrary {
+                           offsets: [
+                               Size(1 bytes),
+                           ],
+                           memory_index: [
+                               0,
+                           ],
+                       },
+                       variants: Single {
+                           index: 1,
+                       },
+                       abi: Aggregate {
+                           sized: true,
+                       },
+                       largest_niche: Some(
+                           Niche {
+                               offset: Size(1 bytes),
+                               value: Int(
+                                   I8,
+                                   false,
+                               ),
+                               valid_range: 0..=1,
+                           },
+                       ),
+                       align: AbiAndPrefAlign {
+                           abi: Align(1 bytes),
+                           pref: $PREF_ALIGN,
+                       },
+                       size: Size(2 bytes),
+                   },
+               ],
+           },
+           abi: Aggregate {
+               sized: true,
+           },
+           largest_niche: Some(
+               Niche {
+                   offset: Size(0 bytes),
+                   value: Int(
+                       I8,
+                       false,
+                   ),
+                   valid_range: 0..=1,
+               },
+           ),
+           align: AbiAndPrefAlign {
+               abi: Align(4 bytes),
+               pref: $PREF_ALIGN,
+           },
+           size: Size(4 bytes),
+       }
+  --> $DIR/zero-sized-array-enum-niche.rs:13:1
+   |
+LL | type AlignedResult = Result<[u32; 0], bool>;
+   | ^^^^^^^^^^^^^^^^^^
+
+error: layout_of(std::result::Result<[u32; 0], Packed<std::num::NonZeroU16>>) = Layout {
+           fields: Arbitrary {
+               offsets: [
+                   Size(0 bytes),
+               ],
+               memory_index: [
+                   0,
+               ],
+           },
+           variants: Multiple {
+               tag: Initialized {
+                   value: Int(
+                       I8,
+                       false,
+                   ),
+                   valid_range: 0..=1,
+               },
+               tag_encoding: Direct,
+               tag_field: 0,
+               variants: [
+                   Layout {
+                       fields: Arbitrary {
+                           offsets: [
+                               Size(4 bytes),
+                           ],
+                           memory_index: [
+                               0,
+                           ],
+                       },
+                       variants: Single {
+                           index: 0,
+                       },
+                       abi: Aggregate {
+                           sized: true,
+                       },
+                       largest_niche: None,
+                       align: AbiAndPrefAlign {
+                           abi: Align(4 bytes),
+                           pref: $PREF_ALIGN,
+                       },
+                       size: Size(4 bytes),
+                   },
+                   Layout {
+                       fields: Arbitrary {
+                           offsets: [
+                               Size(1 bytes),
+                           ],
+                           memory_index: [
+                               0,
+                           ],
+                       },
+                       variants: Single {
+                           index: 1,
+                       },
+                       abi: Aggregate {
+                           sized: true,
+                       },
+                       largest_niche: Some(
+                           Niche {
+                               offset: Size(1 bytes),
+                               value: Int(
+                                   I16,
+                                   false,
+                               ),
+                               valid_range: 1..=65535,
+                           },
+                       ),
+                       align: AbiAndPrefAlign {
+                           abi: Align(1 bytes),
+                           pref: $PREF_ALIGN,
+                       },
+                       size: Size(3 bytes),
+                   },
+               ],
+           },
+           abi: Aggregate {
+               sized: true,
+           },
+           largest_niche: Some(
+               Niche {
+                   offset: Size(0 bytes),
+                   value: Int(
+                       I8,
+                       false,
+                   ),
+                   valid_range: 0..=1,
+               },
+           ),
+           align: AbiAndPrefAlign {
+               abi: Align(4 bytes),
+               pref: $PREF_ALIGN,
+           },
+           size: Size(4 bytes),
+       }
+  --> $DIR/zero-sized-array-enum-niche.rs:29:1
+   |
+LL | type NicheLosesToTagged = Result<[u32; 0], Packed<std::num::NonZeroU16>>;
+   | ^^^^^^^^^^^^^^^^^^^^^^^
+
+error: layout_of(std::result::Result<[u32; 0], Packed<U16IsZero>>) = Layout {
+           fields: Arbitrary {
+               offsets: [
+                   Size(0 bytes),
+               ],
+               memory_index: [
+                   0,
+               ],
+           },
+           variants: Multiple {
+               tag: Initialized {
+                   value: Int(
+                       I16,
+                       false,
+                   ),
+                   valid_range: 0..=1,
+               },
+               tag_encoding: Niche {
+                   dataful_variant: 1,
+                   niche_variants: 0..=0,
+                   niche_start: 1,
+               },
+               tag_field: 0,
+               variants: [
+                   Layout {
+                       fields: Arbitrary {
+                           offsets: [
+                               Size(0 bytes),
+                           ],
+                           memory_index: [
+                               0,
+                           ],
+                       },
+                       variants: Single {
+                           index: 0,
+                       },
+                       abi: Aggregate {
+                           sized: true,
+                       },
+                       largest_niche: None,
+                       align: AbiAndPrefAlign {
+                           abi: Align(4 bytes),
+                           pref: $PREF_ALIGN,
+                       },
+                       size: Size(0 bytes),
+                   },
+                   Layout {
+                       fields: Arbitrary {
+                           offsets: [
+                               Size(0 bytes),
+                           ],
+                           memory_index: [
+                               0,
+                           ],
+                       },
+                       variants: Single {
+                           index: 1,
+                       },
+                       abi: Aggregate {
+                           sized: true,
+                       },
+                       largest_niche: Some(
+                           Niche {
+                               offset: Size(0 bytes),
+                               value: Int(
+                                   I16,
+                                   false,
+                               ),
+                               valid_range: 0..=0,
+                           },
+                       ),
+                       align: AbiAndPrefAlign {
+                           abi: Align(1 bytes),
+                           pref: $PREF_ALIGN,
+                       },
+                       size: Size(2 bytes),
+                   },
+               ],
+           },
+           abi: Aggregate {
+               sized: true,
+           },
+           largest_niche: Some(
+               Niche {
+                   offset: Size(0 bytes),
+                   value: Int(
+                       I16,
+                       false,
+                   ),
+                   valid_range: 0..=1,
+               },
+           ),
+           align: AbiAndPrefAlign {
+               abi: Align(4 bytes),
+               pref: $PREF_ALIGN,
+           },
+           size: Size(4 bytes),
+       }
+  --> $DIR/zero-sized-array-enum-niche.rs:36:1
+   |
+LL | type NicheWinsOverTagged = Result<[u32; 0], Packed<U16IsZero>>;
+   | ^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: aborting due to 3 previous errors
+