fix: Organises the enum niche-filling variant packing implementation.

Organises the niche-repacking solution structure for better mantainability.

Minor bug fix cause by type mismatch.

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

@@ -608,7 +608,7 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
             let mut combined_seed = repr.field_shuffle_seed;
 
             let mut variants_info = IndexVec::<VariantIdx, _>::with_capacity(variants.len());
-            let mut variant_layouts = variants
+            let variant_layouts = variants
                 .iter()
                 .map(|v| {
                     let st = self.univariant(v, repr, StructKind::AlwaysSized).ok()?;
@@ -624,195 +624,27 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                 })
                 .collect::<Option<IndexVec<VariantIdx, _>>>()?;
 
-            let largest_variant_index = variant_layouts
-                .iter_enumerated()
-                .max_by_key(|(_i, layout)| layout.size.bytes())
-                .map(|(i, _layout)| i)?;
-
-            let all_indices = variants.indices();
-            let needs_disc =
-                |index: VariantIdx| index != largest_variant_index && !absent(&variants[index]);
-            let niche_variants = RangeInclusive {
-                start: all_indices.clone().find(|v| needs_disc(*v)).unwrap(),
-                last: all_indices.rev().find(|v| needs_disc(*v)).unwrap(),
-            };
-
-            let count =
-                (niche_variants.last.index() as u128 - niche_variants.start.index() as u128) + 1;
-
-            // Use the largest niche in the largest variant.
-            let niche = variant_layouts[largest_variant_index].largest_niche?;
-            let (niche_start, niche_scalar) = niche.reserve(dl, count)?;
-            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 {
-                    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 {
-                return None;
-            }
-
-            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 }
-            };
-
-            let layout = 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,
-            };
-
-            Some(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 try_niche_variant =
+                |largest_variant_index: VariantIdx| -> Option<LayoutData<FieldIdx, VariantIdx>> {
+                    //so niche_variant candidates don't contaminate each other if one fails
                     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 = RangeInclusive {
+                        start: all_indices.clone().find(|v| needs_disc(*v)).unwrap(),
+                        last: all_indices.rev().find(|v| needs_disc(*v)).unwrap(),
                     };
-                    let niche_variants = niche_variants_start..=niche_variants_end;
 
-                    let count = (niche_variants.end().index() as u128
-                        - niche_variants.start().index() as u128)
+                    let count = (niche_variants.last.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 = variant_layouts[largest_variant_index].largest_niche?;
+                    let (niche_start, niche_scalar) = niche.reserve(dl, count)?;
                     let niche_offset = niche.offset;
                     let niche_size = niche.value.size(dl);
                     let size = variant_layouts[largest_variant_index].size.align_to(align);
@@ -834,41 +666,46 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                             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;
+                            if this_offset + layout.size <= size {
+                                // It'll fit, but we need to make some adjustments.
+                                for offset in layout.field_offsets.iter_mut() {
+                                    *offset += this_offset;
                                 }
-                                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;
 
-                            // 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 };
+                                return true;
+                            }
+                            // Repacking is currently only on future edition. For editions where it is disabled,
+                            // fall back to the original niche-filling behaviour.
+                            if !repr.can_repack_variant_around_niche() {
+                                return false;
+                            }
+                            // 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;
                             }
-                            layout.size += this_offset;
 
-                            true
+                            false
                         });
 
                     if !all_variants_fit {
-                        continue;
+                        return None;
                     }
 
                     let largest_niche = Niche::from_scalar(dl, niche_offset, niche_scalar);
@@ -900,7 +737,7 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                         BackendRepr::Memory { sized: true }
                     };
 
-                    return Some(LayoutData {
+                    let layout = LayoutData {
                         variants: Variants::Multiple {
                             tag: niche_scalar,
                             tag_encoding: TagEncoding::Niche {
@@ -923,18 +760,35 @@ impl<Cx: HasDataLayout> LayoutCalculator<Cx> {
                         max_repr_align,
                         unadjusted_abi_align,
                         randomization_seed: combined_seed,
-                    });
+                    };
+
+                    Some(layout)
+                };
+            if repr.can_repack_variant_around_niche() {
+                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()
+                    })
+                {
+                    if let Some(layout) = try_niche_variant(largest_variant_index) {
+                        return Some(layout);
+                    }
                 }
 
                 None
-            };
-
-        let niche_filling_layout = if repr.can_repack_variant_around_niche() {
-            calculate_niche_filling_layout_repacked()
-        } else {
-            calculate_niche_filling_layout()
+            } else {
+                let largest_variant_index = variant_layouts
+                    .iter_enumerated()
+                    .max_by_key(|(_i, layout)| layout.size.bytes())
+                    .map(|(i, _layout)| i)?;
+                try_niche_variant(largest_variant_index)
+            }
         };
 
+        let niche_filling_layout = calculate_niche_filling_layout();
+
         let discr_type = repr.discr_type();
         let discr_int = Integer::from_attr(dl, discr_type);
         // Because we can only represent one range of valid values, we'll look for the