Do not store tag for single-inhabited-variant enums (if this results in a better layout)

Author: zachs18

compiler/rustc_abi/src/layout.rs

@@ -783,6 +783,9 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
         });
         trace!(?largest_niche);
 
+        let single_variant_layout_eligible =
+            !repr.inhibit_enum_layout_opt() && valid_discriminants.len() == 1;
+
         // `max` is the last valid discriminant before the largest niche
         // `min` is the first valid discriminant after the largest niche
         let (max, min) = largest_niche
@@ -815,10 +818,15 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
         }
 
         // Create the set of structs that represent each variant.
+        let mut single_inhabited_variant_no_tag_layout = None;
         let mut layout_variants = variants
             .iter_enumerated()
             .map(|(i, field_layouts)| {
                 let uninhabited = field_layouts.iter().any(|f| f.is_uninhabited());
+                if !uninhabited && single_variant_layout_eligible {
+                    single_inhabited_variant_no_tag_layout =
+                        Some((i, self.univariant(field_layouts, repr, StructKind::AlwaysSized)));
+                }
                 // We don't need to encode the tag in uninhabited variants in repr(Rust) enums
                 let struct_kind = if uninhabited && !repr.inhibit_enum_layout_opt() {
                     StructKind::AlwaysSized
@@ -845,6 +853,62 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
             })
             .collect::<Result<IndexVec<VariantIdx, _>, _>>()?;
 
+        // If there is a single uninhabited variant, we can use it mostly unchanged as the layout,
+        // without using a tag or niche.
+        //
+        // We do still need to modify it to make all the uninhabited variants fit so they
+        // can be partially-initialized.
+        //
+        // We keep this as a prospective layout, and don't assume it's better than the tagged
+        // layout and return it immediately; e.g. it's worse for `enum Foo { A, B(i32, !) }`
+        // because it has no niche.
+        let no_tag_layout = if let Some((single_inhabited_variant_idx, Ok(mut st))) =
+            single_inhabited_variant_no_tag_layout
+        {
+            // Keep track of original variant layouts (including the inhabited one)
+            // for `offset_of!`.
+            let mut variants = layout_variants.clone();
+            variants[single_inhabited_variant_idx] = st.clone();
+
+            // We know that every other variant is uninhabited, and thus does not have a
+            // prefix for the tag, so we can use them to find the necessary size.
+            for (idx, layout) in layout_variants.iter_enumerated() {
+                if idx != single_inhabited_variant_idx {
+                    st.size = cmp::max(st.size, layout.size);
+                    st.align = st.align.max(layout.align);
+                    st.max_repr_align = st.max_repr_align.max(layout.max_repr_align);
+                    st.unadjusted_abi_align =
+                        st.unadjusted_abi_align.max(layout.unadjusted_abi_align);
+                }
+            }
+
+            // Align the maximum variant size to the largest alignment.
+            st.size = st.size.align_to(st.align.abi);
+
+            // If the inhabited variant's layout would use a non-Memory BackendRepr,
+            // but we made the enum layout bigger or more-aligned due to uninhabited variants,
+            // force the enum to be BackendRepr::Memory.
+            //
+            // FIXME: does this need to care about `max_repr_align` and `unadjusted_abi_align`?
+            // The untagged layout is only used for `repr(Rust)` enums,
+            // so `max_repr_align` and `unadjusted_abi_align` might be irrelevant.
+            if st.size != variants[single_inhabited_variant_idx].size
+                || st.align != variants[single_inhabited_variant_idx].align
+                || st.max_repr_align != variants[single_inhabited_variant_idx].max_repr_align
+                || st.unadjusted_abi_align
+                    != variants[single_inhabited_variant_idx].unadjusted_abi_align
+            {
+                st.backend_repr = BackendRepr::Memory { sized: true };
+            }
+
+            st.variants =
+                Variants::Single { index: single_inhabited_variant_idx, variants: Some(variants) };
+
+            Some(st)
+        } else {
+            None
+        };
+
         // Align the maximum variant size to the largest alignment.
         size = size.align_to(align);
 
@@ -1125,22 +1189,36 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
             randomization_seed: combined_seed,
         };
 
-        let best_layout = match (tagged_layout, niche_filling_layout) {
-            (tl, Some(nl)) => {
-                // Pick the smaller layout; otherwise,
-                // pick the layout with the larger niche; otherwise,
-                // pick tagged as it has simpler codegen.
+        // Pick the smallest layout; otherwise,
+        // pick the layout with the largest niche; otherwise,
+        // pick no_tag as it has simpler codegen than tagged and niched; otherwise,
+        // pick tagged as it has simpler codegen than niched.
+
+        let better_layout_or_first =
+            |l1: LayoutData<FieldIdx, VariantIdx>, l2: LayoutData<FieldIdx, VariantIdx>| {
                 use cmp::Ordering::*;
                 let niche_size = |l: &LayoutData<FieldIdx, VariantIdx>| {
                     l.largest_niche.map_or(0, |n| n.available(dl))
                 };
-                match (tl.size.cmp(&nl.size), niche_size(&tl).cmp(&niche_size(&nl))) {
-                    (Greater, _) => nl,
-                    (Equal, Less) => nl,
-                    _ => tl,
+                match (l1.size.cmp(&l2.size), niche_size(&l1).cmp(&niche_size(&l2))) {
+                    (Greater, _) => l2,
+                    (Equal, Less) => l2,
+                    _ => l1,
                 }
-            }
-            (tl, None) => tl,
+            };
+
+        let best_layout = match niche_filling_layout {
+            None => tagged_layout,
+            // Prefer tagged over niched if they have the same size and niche size,
+            // as the tagged layout has simpler codegen.
+            Some(niched_layout) => better_layout_or_first(tagged_layout, niched_layout),
+        };
+
+        let best_layout = match no_tag_layout {
+            None => best_layout,
+            // Prefer no-tag over tagged/niched if they have the same size and niche size,
+            // as the no-tag layout has simpler codegen.
+            Some(no_tag_layout) => better_layout_or_first(no_tag_layout, best_layout),
         };
 
         Ok(best_layout)

tests/ui/consts/const-eval/ub-enum.rs

@@ -102,8 +102,8 @@ const BAD_UNINHABITED_WITH_DATA2: Result<(i32, !), (i32, Never)> = unsafe { mem:
 
 // All variants have same-size data but only one inhabited.
 // Use `0` as constant to make behavior endianness-independent.
-const GOOD_SEMIINHABITED_WITH_DATA1: Result<i32, (i32, !)> = unsafe { mem::transmute(0u64) };
-const GOOD_SEMIINHABITED_WITH_DATA2: Result<(i32, !), i32> = unsafe { mem::transmute(1u64) };
+const GOOD_SEMIINHABITED_WITH_DATA1: Result<i32, (i32, !)> = unsafe { mem::transmute(0u32) };
+const GOOD_SEMIINHABITED_WITH_DATA2: Result<(i32, !), i32> = unsafe { mem::transmute(0u32) };
 
 const TEST_ICE_89765: () = {
     // This is a regression test for https://github.com/rust-lang/rust/issues/89765.

tests/ui/layout/enum.rs

@@ -35,8 +35,8 @@ enum AbsentVariantUntagged { //~ERROR: size: Size(4 bytes)
 // Even if uninhabited variants are not absent, the enum can still be laid out without
 // a tag.
 #[rustc_layout(size, abi)]
-enum UninhabitedVariantUntagged { //~ERROR: size: Size(8 bytes)
-    //~^ ERROR: abi: ScalarPair(Initialized
+enum UninhabitedVariantUntagged { //~ERROR: size: Size(4 bytes)
+    //~^ ERROR: abi: Scalar(Initialized
     A(i32),
     B(i32, !),
 }
@@ -53,8 +53,8 @@ enum UninhabitedVariantUntaggedBigger { //~ERROR: size: Size(8 bytes)
 }
 
 #[rustc_layout(size, abi)]
-enum UninhabitedVariantWithNiche { //~ERROR: size: Size(3 bytes)
-    //~^ERROR: abi: Memory
+enum UninhabitedVariantWithNiche { //~ERROR: size: Size(2 bytes)
+    //~^ERROR: abi: ScalarPair(Initialized { value: Int(I8, false), valid_range: 0..=1 }, Initialized { value: Int(I8, true), valid_range: 0..=255 })
     A(i8, bool),
     B(u8, u8, !),
 }

tests/ui/layout/enum.stderr

@@ -28,13 +28,13 @@ error: abi: Scalar(Initialized { value: Int(I32, true), valid_range: 0..=4294967
 LL | enum AbsentVariantUntagged {
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-error: size: Size(8 bytes)
+error: size: Size(4 bytes)
   --> $DIR/enum.rs:38:1
    |
 LL | enum UninhabitedVariantUntagged {
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-error: abi: ScalarPair(Initialized { value: Int(I32, false), valid_range: 0..=0 }, Initialized { value: Int(I32, true), valid_range: 0..=4294967295 })
+error: abi: Scalar(Initialized { value: Int(I32, true), valid_range: 0..=4294967295 })
   --> $DIR/enum.rs:38:1
    |
 LL | enum UninhabitedVariantUntagged {
@@ -52,13 +52,13 @@ error: abi: ScalarPair(Initialized { value: Int(I8, false), valid_range: 0..=0 }
 LL | enum UninhabitedVariantUntaggedBigger {
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-error: size: Size(3 bytes)
+error: size: Size(2 bytes)
   --> $DIR/enum.rs:56:1
    |
 LL | enum UninhabitedVariantWithNiche {
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-error: abi: Memory { sized: true }
+error: abi: ScalarPair(Initialized { value: Int(I8, false), valid_range: 0..=1 }, Initialized { value: Int(I8, true), valid_range: 0..=255 })
   --> $DIR/enum.rs:56:1
    |
 LL | enum UninhabitedVariantWithNiche {