Compare tagged/niche-filling layout and pick the best one

src/librustc_middle/lib.rs

@@ -27,6 +27,7 @@
 #![feature(bool_to_option)]
 #![feature(box_patterns)]
 #![feature(box_syntax)]
+#![feature(cmp_min_max_by)]
 #![feature(const_fn)]
 #![feature(const_panic)]
 #![feature(const_fn_transmute)]

src/librustc_middle/ty/layout.rs

@@ -876,6 +876,8 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
                     .iter_enumerated()
                     .all(|(i, v)| v.discr == ty::VariantDiscr::Relative(i.as_u32()));
 
+                let mut niche_filling_layout = None;
+
                 // Niche-filling enum optimization.
                 if !def.repr.inhibit_enum_layout_opt() && no_explicit_discriminants {
                     let mut dataful_variant = None;
@@ -972,7 +974,7 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
                             let largest_niche =
                                 Niche::from_scalar(dl, offset, niche_scalar.clone());
 
-                            return Ok(tcx.intern_layout(Layout {
+                            niche_filling_layout = Some(Layout {
                                 variants: Variants::Multiple {
                                     tag: niche_scalar,
                                     tag_encoding: TagEncoding::Niche {
@@ -991,7 +993,7 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
                                 largest_niche,
                                 size,
                                 align,
-                            }));
+                            });
                         }
                     }
                 }
@@ -1214,7 +1216,7 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
 
                 let largest_niche = Niche::from_scalar(dl, Size::ZERO, tag.clone());
 
-                tcx.intern_layout(Layout {
+                let tagged_layout = Layout {
                     variants: Variants::Multiple {
                         tag,
                         tag_encoding: TagEncoding::Direct,
@@ -1229,7 +1231,23 @@ impl<'tcx> LayoutCx<'tcx, TyCtxt<'tcx>> {
                     abi,
                     align,
                     size,
-                })
+                };
+
+                let best_layout = match (tagged_layout, niche_filling_layout) {
+                    (tagged_layout, Some(niche_filling_layout)) => {
+                        // Pick the smaller layout; otherwise,
+                        // pick the layout with the larger niche; otherwise,
+                        // pick tagged as it has simpler codegen.
+                        cmp::min_by_key(tagged_layout, niche_filling_layout, |layout| {
+                            let niche_size =
+                                layout.largest_niche.as_ref().map_or(0, |n| n.available(dl));
+                            (layout.size, cmp::Reverse(niche_size))
+                        })
+                    }
+                    (tagged_layout, None) => tagged_layout,
+                };
+
+                tcx.intern_layout(best_layout)
             }
 
             // Types with no meaningful known layout.

src/test/ui/print_type_sizes/niche-filling.stdout

@@ -8,12 +8,12 @@ print-type-size     variant `Some`: 12 bytes
 print-type-size         field `.0`: 12 bytes
 print-type-size     variant `None`: 0 bytes
 print-type-size type: `EmbeddedDiscr`: 8 bytes, alignment: 4 bytes
+print-type-size     discriminant: 1 bytes
 print-type-size     variant `Record`: 7 bytes
-print-type-size         field `.val`: 4 bytes
-print-type-size         field `.post`: 2 bytes
 print-type-size         field `.pre`: 1 bytes
+print-type-size         field `.post`: 2 bytes
+print-type-size         field `.val`: 4 bytes
 print-type-size     variant `None`: 0 bytes
-print-type-size     end padding: 1 bytes
 print-type-size type: `MyOption<Union1<std::num::NonZeroU32>>`: 8 bytes, alignment: 4 bytes
 print-type-size     discriminant: 4 bytes
 print-type-size     variant `Some`: 4 bytes

src/test/ui/type-sizes.rs

@@ -5,6 +5,7 @@
 #![feature(never_type)]
 
 use std::mem::size_of;
+use std::num::NonZeroU8;
 
 struct t {a: u8, b: i8}
 struct u {a: u8, b: i8, c: u8}
@@ -102,6 +103,23 @@ enum Option2<A, B> {
     None
 }
 
+// Two layouts are considered for `CanBeNicheFilledButShouldnt`:
+//   Niche-filling:
+//     { u32 (4 bytes), NonZeroU8 + tag in niche (1 byte), padding (3 bytes) }
+//   Tagged:
+//     { tag (1 byte), NonZeroU8 (1 byte), padding (2 bytes), u32 (4 bytes) }
+// Both are the same size (due to padding),
+// but the tagged layout is better as the tag creates a niche with 254 invalid values,
+// allowing types like `Option<Option<CanBeNicheFilledButShouldnt>>` to fit into 8 bytes.
+pub enum CanBeNicheFilledButShouldnt {
+    A(NonZeroU8, u32),
+    B
+}
+pub enum AlwaysTaggedBecauseItHasNoNiche {
+    A(u8, u32),
+    B
+}
+
 pub fn main() {
     assert_eq!(size_of::<u8>(), 1 as usize);
     assert_eq!(size_of::<u32>(), 4 as usize);
@@ -145,4 +163,11 @@ pub fn main() {
     assert_eq!(size_of::<Option<Option<(&(), bool)>>>(), size_of::<(bool, &())>());
     assert_eq!(size_of::<Option<Option2<bool, &()>>>(), size_of::<(bool, &())>());
     assert_eq!(size_of::<Option<Option2<&(), bool>>>(), size_of::<(bool, &())>());
+
+    assert_eq!(size_of::<CanBeNicheFilledButShouldnt>(), 8);
+    assert_eq!(size_of::<Option<CanBeNicheFilledButShouldnt>>(), 8);
+    assert_eq!(size_of::<Option<Option<CanBeNicheFilledButShouldnt>>>(), 8);
+    assert_eq!(size_of::<AlwaysTaggedBecauseItHasNoNiche>(), 8);
+    assert_eq!(size_of::<Option<AlwaysTaggedBecauseItHasNoNiche>>(), 8);
+    assert_eq!(size_of::<Option<Option<AlwaysTaggedBecauseItHasNoNiche>>>(), 8);
 }