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

Author: zachs18

compiler/rustc_abi/src/layout.rs

@@ -197,6 +197,54 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
         })
     }
 
+    pub fn simd_type_for_scalar<FieldIdx: Idx, VariantIdx: Idx, F>(
+        &self,
+        element: Scalar,
+        count: u64,
+        repr_packed: bool,
+    ) -> LayoutCalculatorResult<FieldIdx, VariantIdx, F> {
+        let elt = element;
+        if count == 0 {
+            return Err(LayoutCalculatorError::ZeroLengthSimdType);
+        } else if count > crate::MAX_SIMD_LANES {
+            return Err(LayoutCalculatorError::OversizedSimdType {
+                max_lanes: crate::MAX_SIMD_LANES,
+            });
+        }
+
+        // Compute the size and alignment of the vector
+        let dl = self.cx.data_layout();
+        let size = elt
+            .size(&self.cx)
+            .checked_mul(count, dl)
+            .ok_or_else(|| LayoutCalculatorError::SizeOverflow)?;
+        let (repr, align) = if repr_packed && !count.is_power_of_two() {
+            // Non-power-of-two vectors have padding up to the next power-of-two.
+            // If we're a packed repr, remove the padding while keeping the alignment as close
+            // to a vector as possible.
+            (BackendRepr::Memory { sized: true }, AbiAlign { abi: Align::max_aligned_factor(size) })
+        } else {
+            (BackendRepr::SimdVector { element, count }, dl.llvmlike_vector_align(size))
+        };
+        let size = size.align_to(align.abi);
+
+        Ok(LayoutData {
+            variants: Variants::Single { index: VariantIdx::new(0), variants: None },
+            fields: FieldsShape::Arbitrary {
+                offsets: [Size::ZERO].into(),
+                memory_index: [0].into(),
+            },
+            backend_repr: repr,
+            largest_niche: None,
+            uninhabited: false,
+            size,
+            align,
+            max_repr_align: None,
+            unadjusted_abi_align: elt.align(&self.cx).abi,
+            randomization_seed: (Hash64::new(count)),
+        })
+    }
+
     /// Compute the layout for a coroutine.
     ///
     /// This uses dedicated code instead of [`Self::layout_of_struct_or_enum`], as coroutine
@@ -809,6 +857,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
@@ -841,10 +892,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
@@ -871,6 +927,72 @@ 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.
+        //
+        // FIXME: We shouldn't assume this is better than the tagged layout; it's worse for
+        // `enum Foo { A, B(i32, !) }` because it has no niche.
+        let no_tag_layout = if single_variant_layout_eligible
+            && 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 it bigger or more-aligned due to uninhabited variants,
+            // force it to be BackendRepr::Memory
+            match st.backend_repr {
+                BackendRepr::Scalar(..) | BackendRepr::ScalarPair(..) => {
+                    if st.backend_repr.scalar_size(&self.cx) != Some(st.size)
+                        || st.backend_repr.scalar_align(&self.cx) != Some(st.align.abi)
+                    {
+                        st.backend_repr = BackendRepr::Memory { sized: true }
+                    }
+                }
+                BackendRepr::SimdVector { element, count } => {
+                    // FIXME: is there a better way to do this than making a copy of
+                    // `LayoutCalculator::simd_type` *just* for this?
+                    let vector_layout = self.simd_type_for_scalar::<FieldIdx, VariantIdx, _>(
+                        element,
+                        count,
+                        repr.packed(),
+                    )?;
+                    if vector_layout.size != st.size || vector_layout.align != st.align {
+                        st.backend_repr = BackendRepr::Memory { sized: true }
+                    }
+                }
+                BackendRepr::Memory { .. } => {}
+            }
+
+            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.abi);
 
@@ -1151,22 +1273,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)