feat: Improves enum niche-filling variant packing

This feature only works with "-Zunstable-options --edition future"
as it changes the rustc_abi.

When a variant cannot fit as a whole before or after the reserved niche,
try repacking its fields around the niche instead.

This makes niche-filling less conservative while preserving the existing
layout behavior when whole-variant placement already succeeds.

Signed-off-by: Miguel Marques <miguel.m.marques@tecnico.ulisboa.pt>
Co-authored-by: Vicente Gusmão <vicente.gusmao@tecnico.ulisboa.pt>

Author: Pinguimmar

compiler/rustc_abi/src/layout.rs

@@ -742,7 +742,198 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
             Some(layout)
         };
 
-        let niche_filling_layout = calculate_niche_filling_layout();
+        let calculate_niche_filling_layout_repacked =
+            || -> Option<LayoutData<FieldIdx, VariantIdx>> {
+                struct VariantLayoutInfo {
+                    align_abi: Align,
+                }
+
+                if repr.inhibit_enum_layout_opt() {
+                    return None;
+                }
+
+                if variants.len() < 2 {
+                    return None;
+                }
+
+                let mut align = dl.aggregate_align;
+                let mut max_repr_align = repr.align;
+                let mut unadjusted_abi_align = align;
+                let mut combined_seed = repr.field_shuffle_seed;
+
+                let mut variants_info = IndexVec::<VariantIdx, _>::with_capacity(variants.len());
+                let variant_layouts = variants
+                    .iter()
+                    .map(|v| {
+                        let st = self.univariant(v, repr, StructKind::AlwaysSized).ok()?;
+
+                        variants_info.push(VariantLayoutInfo { align_abi: st.align.abi });
+
+                        align = align.max(st.align.abi);
+                        max_repr_align = max_repr_align.max(st.max_repr_align);
+                        unadjusted_abi_align = unadjusted_abi_align.max(st.unadjusted_abi_align);
+                        combined_seed = combined_seed.wrapping_add(st.randomization_seed);
+
+                        Some(VariantLayout::from_layout(st))
+                    })
+                    .collect::<Option<IndexVec<VariantIdx, _>>>()?;
+
+                let max_variant_size = variant_layouts.iter().map(|layout| layout.size).max()?;
+
+                // Chooses the first max-sized niche-providing variant whose layout succeeds.
+                for (largest_variant_index, _) in
+                    variant_layouts.iter_enumerated().filter(|&(_, layout)| {
+                        layout.size == max_variant_size && layout.largest_niche.is_some()
+                    })
+                {
+                    let mut variant_layouts = variant_layouts.clone();
+
+                    let all_indices = variants.indices();
+                    let needs_disc = |index: VariantIdx| {
+                        index != largest_variant_index && !absent(&variants[index])
+                    };
+                    let Some(niche_variants_start) = all_indices.clone().find(|v| needs_disc(*v))
+                    else {
+                        continue;
+                    };
+                    let Some(niche_variants_end) = all_indices.rev().find(|v| needs_disc(*v))
+                    else {
+                        continue;
+                    };
+                    let niche_variants = niche_variants_start..=niche_variants_end;
+
+                    let count = (niche_variants.end().index() as u128
+                        - niche_variants.start().index() as u128)
+                        + 1;
+
+                    // Use the largest niche in the largest variant.
+                    let Some(niche) = variant_layouts[largest_variant_index].largest_niche else {
+                        continue;
+                    };
+                    let Some((niche_start, niche_scalar)) = niche.reserve(dl, count) else {
+                        continue;
+                    };
+                    let niche_offset = niche.offset;
+                    let niche_size = niche.value.size(dl);
+                    let size = variant_layouts[largest_variant_index].size.align_to(align);
+
+                    let all_variants_fit =
+                        variant_layouts.iter_enumerated_mut().all(|(i, layout)| {
+                            if i == largest_variant_index {
+                                return true;
+                            }
+
+                            layout.largest_niche = None;
+
+                            if layout.size <= niche_offset {
+                                // This variant will fit before the niche.
+                                return true;
+                            }
+
+                            // Determine if it'll fit after the niche.
+                            let this_align = variants_info[i].align_abi;
+                            let this_offset = (niche_offset + niche_size).align_to(this_align);
+
+                            if this_offset + layout.size > size {
+                                // The ordinary niche-filling path can only move a non-largest variant as a
+                                // whole before or after the chosen niche. If that fails, try placing the
+                                // variant's fields individually while treating the niche bytes as reserved.
+                                if let Some(repacked) = self.try_layout_variant_around_niche(
+                                    &variants[i],
+                                    repr,
+                                    i,
+                                    size,
+                                    align,
+                                    niche_offset,
+                                    niche_size,
+                                ) {
+                                    *layout = VariantLayout::from_layout(repacked);
+                                    return true;
+                                }
+                                return false;
+                            }
+
+                            // It'll fit, but we need to make some adjustments.
+                            for offset in layout.field_offsets.iter_mut() {
+                                *offset += this_offset;
+                            }
+
+                            // It can't be a Scalar or ScalarPair because the offset isn't 0.
+                            if !layout.is_uninhabited() {
+                                layout.backend_repr = BackendRepr::Memory { sized: true };
+                            }
+                            layout.size += this_offset;
+
+                            true
+                        });
+
+                    if !all_variants_fit {
+                        continue;
+                    }
+
+                    let largest_niche = Niche::from_scalar(dl, niche_offset, niche_scalar);
+
+                    let others_zst = variant_layouts
+                        .iter_enumerated()
+                        .all(|(i, layout)| i == largest_variant_index || layout.size == Size::ZERO);
+                    let same_size = size == variant_layouts[largest_variant_index].size;
+                    let same_align = align == variants_info[largest_variant_index].align_abi;
+
+                    let uninhabited = variant_layouts.iter().all(|v| v.is_uninhabited());
+                    let abi = if same_size && same_align && others_zst {
+                        match variant_layouts[largest_variant_index].backend_repr {
+                            // When the total alignment and size match, we can use the
+                            // same ABI as the scalar variant with the reserved niche.
+                            BackendRepr::Scalar(_) => BackendRepr::Scalar(niche_scalar),
+                            BackendRepr::ScalarPair(first, second) => {
+                                // Only the niche is guaranteed to be initialised,
+                                // so use union layouts for the other primitive.
+                                if niche_offset == Size::ZERO {
+                                    BackendRepr::ScalarPair(niche_scalar, second.to_union())
+                                } else {
+                                    BackendRepr::ScalarPair(first.to_union(), niche_scalar)
+                                }
+                            }
+                            _ => BackendRepr::Memory { sized: true },
+                        }
+                    } else {
+                        BackendRepr::Memory { sized: true }
+                    };
+
+                    return Some(LayoutData {
+                        variants: Variants::Multiple {
+                            tag: niche_scalar,
+                            tag_encoding: TagEncoding::Niche {
+                                untagged_variant: largest_variant_index,
+                                niche_variants,
+                                niche_start,
+                            },
+                            tag_field: FieldIdx::new(0),
+                            variants: variant_layouts,
+                        },
+                        fields: FieldsShape::Arbitrary {
+                            offsets: [niche_offset].into(),
+                            in_memory_order: [FieldIdx::new(0)].into(),
+                        },
+                        backend_repr: abi,
+                        largest_niche,
+                        uninhabited,
+                        size,
+                        align: AbiAlign::new(align),
+                        max_repr_align,
+                        unadjusted_abi_align,
+                        randomization_seed: combined_seed,
+                    });
+                }
+
+                None
+            };
+
+        let niche_filling_layout = if repr.can_repack_variant_around_niche() {
+            calculate_niche_filling_layout_repacked()
+        } else {
+            calculate_niche_filling_layout()
+        };
 
         let discr_type = repr.discr_type();
         let discr_int = Integer::from_attr(dl, discr_type);
@@ -1091,6 +1282,121 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
         Ok(best_layout)
     }
 
+    fn try_layout_variant_around_niche<
+        'a,
+        FieldIdx: Idx,
+        VariantIdx: Idx,
+        F: Deref<Target = &'a LayoutData<FieldIdx, VariantIdx>> + fmt::Debug + Copy,
+    >(
+        &self,
+        fields: &IndexSlice<FieldIdx, F>,
+        repr: &ReprOptions,
+        variant_index: VariantIdx,
+        total_size: Size,
+        total_align: Align,
+        niche_offset: Size,
+        niche_size: Size,
+    ) -> Option<LayoutData<FieldIdx, VariantIdx>> {
+        let dl = self.cx.data_layout();
+
+        if repr.inhibit_struct_field_reordering() {
+            return None;
+        }
+        if fields.iter().any(|f| f.is_unsized()) {
+            return None;
+        }
+
+        // This is only a necessary condition: alignment can still make placement fail below.
+        let min_used_size = fields.iter().try_fold(niche_size, |size, field| {
+            if field.is_zst() { Some(size) } else { size.checked_add(field.size, dl) }
+        })?;
+        if min_used_size > total_size {
+            return None;
+        }
+
+        let mut offsets = IndexVec::from_elem(Size::ZERO, fields);
+        let mut in_memory_order: IndexVec<u32, FieldIdx> = fields.indices().collect();
+
+        in_memory_order.raw.sort_by_key(|&i| {
+            let field = &fields[i];
+
+            let field_align = if let Some(pack) = repr.pack {
+                field.align.min(AbiAlign::new(pack))
+            } else {
+                field.align
+            };
+
+            (cmp::Reverse(field_align.abi.bytes()), cmp::Reverse(field.size.bytes()))
+        });
+
+        let niche_end = niche_offset.checked_add(niche_size, dl)?;
+        let mut occupied = vec![(niche_offset, niche_end)];
+
+        let mut max_repr_align = repr.align;
+        let mut unadjusted_abi_align =
+            if repr.pack.is_some() { dl.i8_align } else { dl.aggregate_align };
+
+        for &i in &in_memory_order {
+            let field = &fields[i];
+
+            let field_align = if let Some(pack) = repr.pack {
+                field.align.min(AbiAlign::new(pack))
+            } else {
+                field.align
+            };
+
+            max_repr_align = max_repr_align.max(field.max_repr_align);
+            unadjusted_abi_align = unadjusted_abi_align.max(field_align.abi);
+
+            let mut candidate = Size::ZERO;
+            'search: loop {
+                candidate = candidate.align_to(field_align.abi);
+
+                let end = candidate.checked_add(field.size, dl)?;
+                if end > total_size {
+                    return None;
+                }
+
+                for &(occupied_start, occupied_end) in &occupied {
+                    if end <= occupied_start {
+                        break;
+                    }
+                    if candidate < occupied_end {
+                        candidate = occupied_end;
+                        continue 'search;
+                    }
+                }
+
+                offsets[i] = candidate;
+
+                if !field.is_zst() {
+                    occupied.push((candidate, end));
+                    occupied.sort_by_key(|&(start, _end)| start);
+                }
+
+                break;
+            }
+        }
+
+        in_memory_order.raw.sort_by_key(|&i| offsets[i]);
+
+        Some(LayoutData {
+            variants: Variants::Single { index: variant_index },
+            fields: FieldsShape::Arbitrary { offsets, in_memory_order },
+            backend_repr: BackendRepr::Memory { sized: true },
+            largest_niche: None,
+            uninhabited: fields.iter().any(|f| f.is_uninhabited()),
+            align: AbiAlign::new(total_align),
+            size: total_size,
+            max_repr_align,
+            unadjusted_abi_align,
+            randomization_seed: fields
+                .iter()
+                .fold(Hash64::ZERO, |acc, f| acc.wrapping_add(f.randomization_seed))
+                .wrapping_add(repr.field_shuffle_seed),
+        })
+    }
+
     fn univariant_biased<
         'a,
         FieldIdx: Idx,

compiler/rustc_abi/src/lib.rs

@@ -94,7 +94,9 @@ bitflags! {
         /// See [`TyAndLayout::pass_indirectly_in_non_rustic_abis`] for details.
         const PASS_INDIRECTLY_IN_NON_RUSTIC_ABIS = 1 << 5;
         const IS_SCALABLE        = 1 << 6;
-         // Any of these flags being set prevent field reordering optimisation.
+        /// If true, enum niche-filling may repack variant fields around the reserved niche.
+        const CAN_REPACK_VARIANT_AROUND_NICHE = 1 << 7;
+        // Any of these flags being set prevent field reordering optimisation.
         const FIELD_ORDER_UNOPTIMIZABLE = ReprFlags::IS_C.bits()
                                  | ReprFlags::IS_SIMD.bits()
                                  | ReprFlags::IS_SCALABLE.bits()
@@ -226,6 +228,12 @@ impl ReprOptions {
         !self.inhibit_struct_field_reordering() && self.flags.contains(ReprFlags::RANDOMIZE_LAYOUT)
     }
 
+    /// Returns `true` if enum niche-filling can try to pack variant fields around
+    /// the reserved niche when whole-variant placement fails.
+    pub fn can_repack_variant_around_niche(&self) -> bool {
+        self.flags.contains(ReprFlags::CAN_REPACK_VARIANT_AROUND_NICHE)
+    }
+
     /// Returns `true` if this `#[repr()]` should inhibit union ABI optimisations.
     pub fn inhibits_union_abi_opt(&self) -> bool {
         self.c()

compiler/rustc_middle/src/ty/mod.rs

@@ -1603,6 +1603,10 @@ impl<'tcx> TyCtxt<'tcx> {
             flags.insert(ReprFlags::RANDOMIZE_LAYOUT);
         }
 
+        if self.def_span(did).edition().at_least_edition_future() {
+            flags.insert(ReprFlags::CAN_REPACK_VARIANT_AROUND_NICHE);
+        }
+
         // box is special, on the one hand the compiler assumes an ordered layout, with the pointer
         // always at offset zero. On the other hand we want scalar abi optimizations.
         let is_box = self.is_lang_item(did.to_def_id(), LangItem::OwnedBox);