Only use the padding for the discriminent if it is bigger than the biggest niche

Author: ogoffart

src/librustc_middle/ty/layout.rs

@@ -900,6 +900,7 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
                     let mut niche_variants = VariantIdx::MAX..=VariantIdx::new(0);
                     let mut max_size = Size::ZERO;
                     let mut second_max_size = Size::ZERO;
+                    let mut align = dl.aggregate_align;
 
                     // The size computations below assume that the padding is minimum.
                     // This is the case when fields are re-ordered.
@@ -909,7 +910,6 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
                             niche_variants =
                                 *niche_variants.start().min(&d)..=*niche_variants.end().max(&d);
                         };
-                        let mut align = dl.aggregate_align;
 
                         // Find the largest and second largest variant.
                         for (v, fields) in variants.iter_enumerated() {
@@ -939,10 +939,6 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
                                 extend_niche_range(v);
                             }
                         }
-                        if !max_size.is_aligned(align.abi) {
-                            // Don't perform niche optimisation if there is padding anyway
-                            dataful_variant = None;
-                        }
                     } else {
                         // Find one non-ZST variant.
                         'variants: for (v, fields) in variants.iter_enumerated() {
@@ -980,12 +976,22 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
                             .filter_map(|(j, &field)| Some((j, field.largest_niche.as_ref()?)))
                             .max_by_key(|(_, n)| (n.available(dl), cmp::Reverse(n.offset)))
                             .and_then(|(field_index, niche)| {
-                                // make sure there is enough room for the other variants
-                                if struct_reordering_opt
-                                    && max_size - (niche.offset + niche.scalar.value.size(dl))
+                                if struct_reordering_opt {
+                                    // make sure there is enough room for the other variants
+                                    if max_size - (niche.offset + niche.scalar.value.size(dl))
                                         < second_max_size
-                                {
-                                    return None;
+                                    {
+                                        return None;
+                                    }
+                                    // Don't perform niche optimisation if there is more padding than
+                                    // what's available in the biggest niche
+                                    let padding_bits =
+                                        (max_size.align_to(align.abi) - max_size).bits();
+                                    if padding_bits >= 128
+                                        || (1 << padding_bits) > niche.available(dl)
+                                    {
+                                        return None;
+                                    }
                                 }
                                 Some((field_index, niche, niche.reserve(self, count)?))
                             });
@@ -1028,7 +1034,10 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
                             } else {
                                 st[dataful_variant].fields.offset(field_index) + niche.offset
                             };
-                            let size = st[dataful_variant].size;
+                            let size = st[dataful_variant].size.align_to(align.abi);
+                            debug_assert!(
+                                !struct_reordering_opt || size == max_size.align_to(align.abi)
+                            );
                             debug_assert!(st.iter().all(|v| { v.size <= size }));
 
                             let abi = if st.iter().all(|v| v.abi.is_uninhabited()) {

src/test/ui/type-sizes.rs

@@ -112,6 +112,34 @@ enum NicheWithData {
     D(u32, u32),
 }
 
+// A type with almost 2^16 invalid values.
+#[repr(u16)]
+pub enum NicheU16 {
+    _0,
+}
+
+pub enum EnumManyVariant<X> {
+    Dataful(u8, X),
+
+    // 0x100 niche variants.
+    _00, _01, _02, _03, _04, _05, _06, _07, _08, _09, _0A, _0B, _0C, _0D, _0E, _0F,
+    _10, _11, _12, _13, _14, _15, _16, _17, _18, _19, _1A, _1B, _1C, _1D, _1E, _1F,
+    _20, _21, _22, _23, _24, _25, _26, _27, _28, _29, _2A, _2B, _2C, _2D, _2E, _2F,
+    _30, _31, _32, _33, _34, _35, _36, _37, _38, _39, _3A, _3B, _3C, _3D, _3E, _3F,
+    _40, _41, _42, _43, _44, _45, _46, _47, _48, _49, _4A, _4B, _4C, _4D, _4E, _4F,
+    _50, _51, _52, _53, _54, _55, _56, _57, _58, _59, _5A, _5B, _5C, _5D, _5E, _5F,
+    _60, _61, _62, _63, _64, _65, _66, _67, _68, _69, _6A, _6B, _6C, _6D, _6E, _6F,
+    _70, _71, _72, _73, _74, _75, _76, _77, _78, _79, _7A, _7B, _7C, _7D, _7E, _7F,
+    _80, _81, _82, _83, _84, _85, _86, _87, _88, _89, _8A, _8B, _8C, _8D, _8E, _8F,
+    _90, _91, _92, _93, _94, _95, _96, _97, _98, _99, _9A, _9B, _9C, _9D, _9E, _9F,
+    _A0, _A1, _A2, _A3, _A4, _A5, _A6, _A7, _A8, _A9, _AA, _AB, _AC, _AD, _AE, _AF,
+    _B0, _B1, _B2, _B3, _B4, _B5, _B6, _B7, _B8, _B9, _BA, _BB, _BC, _BD, _BE, _BF,
+    _C0, _C1, _C2, _C3, _C4, _C5, _C6, _C7, _C8, _C9, _CA, _CB, _CC, _CD, _CE, _CF,
+    _D0, _D1, _D2, _D3, _D4, _D5, _D6, _D7, _D8, _D9, _DA, _DB, _DC, _DD, _DE, _DF,
+    _E0, _E1, _E2, _E3, _E4, _E5, _E6, _E7, _E8, _E9, _EA, _EB, _EC, _ED, _EE, _EF,
+    _F0, _F1, _F2, _F3, _F4, _F5, _F6, _F7, _F8, _F9, _FA, _FB, _FC, _FD, _FE, _FF,
+}
+
 pub fn main() {
     assert_eq!(size_of::<u8>(), 1 as usize);
     assert_eq!(size_of::<u32>(), 4 as usize);
@@ -172,4 +200,14 @@ pub fn main() {
     assert_eq!(size_of::<FillPadding>(), 8);
     assert_eq!(size_of::<Option<FillPadding>>(), 8);
     assert_eq!(size_of::<Option<Option<FillPadding>>>(), 8);
+
+    assert_eq!(size_of::<Result<(std::num::NonZeroU8, u8, u8), u16>>(), 4);
+    assert_eq!(size_of::<Option<Result<(std::num::NonZeroU8, u8, u8), u16>>>(), 4);
+    assert_eq!(size_of::<Result<(std::num::NonZeroU8, u8, u8, u8), u16>>(), 4);
+
+    assert_eq!(size_of::<EnumManyVariant<u16>>(), 6);
+    assert_eq!(size_of::<EnumManyVariant<NicheU16>>(), 4);
+    assert_eq!(size_of::<EnumManyVariant<Option<NicheU16>>>(), 4);
+    assert_eq!(size_of::<EnumManyVariant<Option2<NicheU16,u8>>>(), 6);
+    assert_eq!(size_of::<EnumManyVariant<Option<(NicheU16,u8)>>>(), 6);
 }