frontend 0.4.1

use alloc::vec::Vec;
use core::fmt;
use core::ops::{Deref, Range};

use crate::rustc_data_structures::intern::Interned;
use crate::rustc_data_structures::range_set::RangeSet;
use rustc_macros::StableHash;

use crate::rustc_abi::layout::{FieldIdx, VariantIdx};
use crate::rustc_abi::{
    AbiAlign, Align, BackendRepr, FieldsShape, Float, HasDataLayout, LayoutData, Niche,
    PointeeInfo, Primitive, Size, Variants,
};

// Explicitly import `Float` to avoid ambiguity with `Primitive::Float`.

#[derive(Copy, Clone, PartialEq, Eq, Hash, StableHash)]
pub struct Layout<'a>(pub Interned<'a, LayoutData<FieldIdx, VariantIdx>>);

impl<'a> fmt::Debug for Layout<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // See comment on `<LayoutData as Debug>::fmt` above.
        self.0.0.fmt(f)
    }
}

impl<'a> Deref for Layout<'a> {
    type Target = &'a LayoutData<FieldIdx, VariantIdx>;
    fn deref(&self) -> &&'a LayoutData<FieldIdx, VariantIdx> {
        &self.0.0
    }
}

impl<'a> Layout<'a> {
    pub fn fields(self) -> &'a FieldsShape<FieldIdx> {
        &self.0.0.fields
    }

    pub fn variants(self) -> &'a Variants<FieldIdx, VariantIdx> {
        &self.0.0.variants
    }

    pub fn backend_repr(self) -> BackendRepr {
        self.0.0.backend_repr
    }

    pub fn largest_niche(self) -> Option<Niche> {
        self.0.0.largest_niche
    }

    pub fn align(self) -> AbiAlign {
        self.0.0.align
    }

    pub fn size(self) -> Size {
        self.0.0.size
    }

    pub fn max_repr_align(self) -> Option<Align> {
        self.0.0.max_repr_align
    }

    pub fn unadjusted_abi_align(self) -> Align {
        self.0.0.unadjusted_abi_align
    }
}

/// The layout of a type, alongside the type itself.
/// Provides various type traversal APIs (e.g., recursing into fields).
///
/// Note that the layout is NOT guaranteed to always be identical
/// to that obtained from `layout_of(ty)`, as we need to produce
/// layouts for which Rust types do not exist, such as enum variants
/// or synthetic fields of enums (i.e., discriminants) and wide pointers.
#[derive(Copy, Clone, PartialEq, Eq, Hash, StableHash)]
pub struct TyAndLayout<'a, Ty> {
    pub ty: Ty,
    pub layout: Layout<'a>,
}

impl<'a, Ty: fmt::Display> fmt::Debug for TyAndLayout<'a, Ty> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Print the type in a readable way, not its debug representation.
        f.debug_struct("TyAndLayout")
            .field("ty", &format_args!("{}", self.ty))
            .field("layout", &self.layout)
            .finish()
    }
}

impl<'a, Ty> Deref for TyAndLayout<'a, Ty> {
    type Target = &'a LayoutData<FieldIdx, VariantIdx>;
    fn deref(&self) -> &&'a LayoutData<FieldIdx, VariantIdx> {
        &self.layout.0.0
    }
}

impl<'a, Ty> AsRef<LayoutData<FieldIdx, VariantIdx>> for TyAndLayout<'a, Ty> {
    fn as_ref(&self) -> &LayoutData<FieldIdx, VariantIdx> {
        &*self.layout.0.0
    }
}

/// Trait that needs to be implemented by the higher-level type representation
/// (e.g. `crate::rustc_middle::ty::Ty`), to provide `crate::rustc_target::abi` functionality.
pub trait TyAbiInterface<'a, C>: Sized + core::fmt::Debug + core::fmt::Display {
    fn ty_and_layout_for_variant(
        this: TyAndLayout<'a, Self>,
        cx: &C,
        variant_index: VariantIdx,
    ) -> TyAndLayout<'a, Self>;
    fn ty_and_layout_field(this: TyAndLayout<'a, Self>, cx: &C, i: usize) -> TyAndLayout<'a, Self>;
    fn ty_and_layout_pointee_info_at(
        this: TyAndLayout<'a, Self>,
        cx: &C,
        offset: Size,
    ) -> Option<PointeeInfo>;
    fn is_adt(this: TyAndLayout<'a, Self>) -> bool;
    fn is_never(this: TyAndLayout<'a, Self>) -> bool;
    fn is_tuple(this: TyAndLayout<'a, Self>) -> bool;
    fn is_unit(this: TyAndLayout<'a, Self>) -> bool;
    fn is_transparent(this: TyAndLayout<'a, Self>) -> bool;
    fn is_complex_number_lang_item(this: TyAndLayout<'a, Self>, cx: &C) -> bool;
    fn is_scalable_vector(this: TyAndLayout<'a, Self>) -> bool;
    /// See [`TyAndLayout::pass_indirectly_in_non_rustic_abis`] for details.
    fn is_pass_indirectly_in_non_rustic_abis_flag_set(this: TyAndLayout<'a, Self>) -> bool;
}

impl<'a, Ty> TyAndLayout<'a, Ty> {
    /// Synthetize a layout representing the variant-specific fields of an enum-like layout.
    ///
    /// Note that the resulting layout *does not* fully describes `self.ty` at that specific
    /// variant: prefix fields (e.g. in coroutines) and tag information are lost.
    ///
    /// If you don't need type information about the variant's fields, prefer using
    /// `self.layout.variants` directly.
    pub fn for_variant<C>(self, cx: &C, variant_index: VariantIdx) -> Self
    where
        Ty: TyAbiInterface<'a, C>,
    {
        Ty::ty_and_layout_for_variant(self, cx, variant_index)
    }

    pub fn field<C>(self, cx: &C, i: usize) -> Self
    where
        Ty: TyAbiInterface<'a, C>,
    {
        Ty::ty_and_layout_field(self, cx, i)
    }

    pub fn pointee_info_at<C>(self, cx: &C, offset: Size) -> Option<PointeeInfo>
    where
        Ty: TyAbiInterface<'a, C>,
    {
        Ty::ty_and_layout_pointee_info_at(self, cx, offset)
    }

    pub fn is_single_fp_element<C>(self, cx: &C) -> bool
    where
        Ty: TyAbiInterface<'a, C>,
        C: HasDataLayout,
    {
        match self.backend_repr {
            BackendRepr::Scalar(scalar) => {
                matches!(scalar.primitive(), Primitive::Float(Float::F32 | Float::F64))
            }
            BackendRepr::Memory { .. } => {
                if self.fields.count() == 1 && self.fields.offset(0).bytes() == 0 {
                    self.field(cx, 0).is_single_fp_element(cx)
                } else {
                    false
                }
            }
            _ => false,
        }
    }

    pub fn is_single_vector_element<C>(self, cx: &C, expected_size: Size) -> bool
    where
        Ty: TyAbiInterface<'a, C>,
        C: HasDataLayout,
    {
        match self.backend_repr {
            BackendRepr::SimdVector { .. } => self.size == expected_size,
            BackendRepr::Memory { .. } => {
                if self.fields.count() == 1 && self.fields.offset(0).bytes() == 0 {
                    self.field(cx, 0).is_single_vector_element(cx, expected_size)
                } else {
                    false
                }
            }
            _ => false,
        }
    }

    pub fn is_adt<C>(self) -> bool
    where
        Ty: TyAbiInterface<'a, C>,
    {
        Ty::is_adt(self)
    }

    pub fn is_never<C>(self) -> bool
    where
        Ty: TyAbiInterface<'a, C>,
    {
        Ty::is_never(self)
    }

    pub fn is_tuple<C>(self) -> bool
    where
        Ty: TyAbiInterface<'a, C>,
    {
        Ty::is_tuple(self)
    }

    pub fn is_unit<C>(self) -> bool
    where
        Ty: TyAbiInterface<'a, C>,
    {
        Ty::is_unit(self)
    }

    pub fn is_transparent<C>(self) -> bool
    where
        Ty: TyAbiInterface<'a, C>,
    {
        Ty::is_transparent(self)
    }

    /// Returns `true` if this type needs to match the ABI of the C `_Complex` type. See
    /// [`TyAndLayout::complex_number_primitive`] for details.
    pub fn is_complex_number<C>(self, cx: &C) -> bool
    where
        Ty: TyAbiInterface<'a, C> + Copy,
    {
        self.complex_number_primitive(cx).is_some()
    }

    pub fn is_scalable_vector<C>(self) -> bool
    where
        Ty: TyAbiInterface<'a, C>,
    {
        Ty::is_scalable_vector(self)
    }

    /// If this method returns `true`, then this type should always have a `PassMode` of
    /// `Indirect { on_stack: false, .. }` when being used as the argument type of a function with a
    /// non-Rustic ABI (this is true for structs annotated with the
    /// `#[rustc_pass_indirectly_in_non_rustic_abis]` attribute).
    ///
    /// This is used to replicate some of the behaviour of C array-to-pointer decay; however unlike
    /// C any changes the caller makes to the passed value will not be reflected in the callee, so
    /// the attribute is only useful for types where observing the value in the caller after the
    /// function call isn't allowed (a.k.a. `va_list`).
    ///
    /// This function handles transparent types automatically.
    pub fn pass_indirectly_in_non_rustic_abis<C>(self, cx: &C) -> bool
    where
        Ty: TyAbiInterface<'a, C> + Copy,
    {
        let base = self.peel_transparent_wrappers(cx);
        Ty::is_pass_indirectly_in_non_rustic_abis_flag_set(base)
    }

    /// Recursively peel away transparent wrappers, returning the inner value.
    ///
    /// The return value is not `repr(transparent)` and/or does
    /// not have a non-1zst field.
    pub fn peel_transparent_wrappers<C>(mut self, cx: &C) -> Self
    where
        Ty: TyAbiInterface<'a, C> + Copy,
    {
        while self.is_transparent()
            && let Some((_, field)) = self.non_1zst_field(cx)
        {
            self = field;
        }

        self
    }

    /// Finds the one field that is not a 1-ZST.
    /// Returns `None` if there are multiple non-1-ZST fields or only 1-ZST-fields.
    pub fn non_1zst_field<C>(&self, cx: &C) -> Option<(FieldIdx, Self)>
    where
        Ty: TyAbiInterface<'a, C> + Copy,
    {
        let mut found = None;
        for field_idx in 0..self.fields.count() {
            let field = self.field(cx, field_idx);
            if field.is_1zst() {
                continue;
            }
            if found.is_some() {
                // More than one non-1-ZST field.
                return None;
            }
            found = Some((FieldIdx::from_usize(field_idx), field));
        }
        found
    }

    /// If this type should match the ABI of the C `_Complex` type, returns the primitive that is
    /// used for its parts. This only returns `Some(T)` for `core::num::Complex<T>` where `T` is
    /// either a float or an integer. `repr(transparent)` wrapper types are automatically handled.
    pub fn complex_number_primitive<C>(&self, cx: &C) -> Option<Primitive>
    where
        Ty: TyAbiInterface<'a, C> + Copy,
    {
        let complex = self.peel_transparent_wrappers(cx);
        if !Ty::is_complex_number_lang_item(complex, cx) {
            return None;
        }

        let part = complex.field(cx, 0).peel_transparent_wrappers(cx);

        if let BackendRepr::Scalar(scalar) = part.backend_repr {
            // Only Complex<{ float }> and Complex<{ integer }> have special layout.
            let primitive = scalar.primitive();
            match primitive {
                // Explicitly spell out all the float types so that any new ones have to be added to
                // one of the match branches.
                Primitive::Int(..)
                | Primitive::Float(Float::F16 | Float::F32 | Float::F64 | Float::F128) => {
                    Some(primitive)
                }
                Primitive::Pointer(..) => None,
            }
        } else {
            None
        }
    }

    /// Returns `Some` if this type has the ABI of the C `_Complex` type with float parts. See
    /// [`TyAndLayout::complex_number_primitive`] for details.
    pub fn complex_float<C>(&self, cx: &C) -> Option<Float>
    where
        Ty: TyAbiInterface<'a, C> + Copy,
    {
        if let Some(Primitive::Float(float)) = self.complex_number_primitive(cx) {
            Some(float)
        } else {
            None
        }
    }

    /// Whether this type/layout has any padding that is dependent on a variant, i.e. has bytes that
    /// are padding for some, but not all, valid values of this type.
    pub fn has_variant_dependent_padding<C>(&self, cx: &C) -> bool
    where
        Ty: TyAbiInterface<'a, C> + Copy,
    {
        match self.variants {
            Variants::Multiple { .. } => true,
            Variants::Empty => false,
            Variants::Single { .. } => match &self.fields {
                FieldsShape::Primitive | FieldsShape::Union(_) => false,
                FieldsShape::Array { count, .. } => {
                    *count > 0 && self.field(cx, 0).has_variant_dependent_padding(cx)
                }
                FieldsShape::Arbitrary { offsets, .. } => {
                    (0..offsets.len()).any(|i| self.field(cx, i).has_variant_dependent_padding(cx))
                }
            },
        }
    }

    /// The ranges of bytes that are always ignored by the representation relation of this type.
    ///
    /// In other words, for any sequence of bytes, if we reset the these padding bytes to uninit,
    /// then these two sequences of bytes represent the same value (or they are both invalid).
    /// This is the "guaranteed" padding. There may be more bytes that are padding for some
    /// but not all variants of this type; those are not included.
    /// (E.g. `Option<i8>` has no guaranteed padding so the empty range set is returned, but its `None` value still has padding).
    pub fn variant_independent_padding_ranges<C>(&self, cx: &C) -> Vec<Range<Size>>
    where
        Ty: TyAbiInterface<'a, C> + Copy,
    {
        let mut data = RangeSet::new();
        self.add_data_ranges(cx, Size::ZERO, &mut data);

        // Find gaps between the data ranges.
        let mut uninit_ranges = Vec::new();
        let mut covered_until = Size::ZERO;
        for &(offset, size) in data.0.iter() {
            if offset > covered_until {
                uninit_ranges.push(covered_until..offset);
            }
            covered_until = Ord::max(covered_until, offset + size);
        }

        // Add trailing padding.
        if self.size > covered_until {
            uninit_ranges.push(covered_until..self.size);
        }

        uninit_ranges
    }

    /// The ranges of bytes that are ignored by the representation relation of this variant.
    ///
    /// The result does not include variant-independent padding.
    pub fn variant_dependent_padding_ranges<C>(
        &self,
        cx: &C,
        variant_index: VariantIdx,
    ) -> Vec<Range<Size>>
    where
        Ty: TyAbiInterface<'a, C> + Copy,
    {
        let Variants::Multiple { .. } = self.variants else {
            return Vec::new();
        };

        // Bytes that are data in some variant.
        let mut any = RangeSet::new();
        self.add_data_ranges(cx, Size::ZERO, &mut any);

        // Bytes that are data in this variant.
        let mut this = RangeSet::new();

        // The variants do not contain e.g. the discriminant or coroutine upvars.
        let FieldsShape::Arbitrary { offsets, in_memory_order: _ } = &self.fields else {
            unreachable!("a multi-variant layout should have `Arbitrary` fields")
        };

        // So add them explicitly.
        for (field, &offset) in offsets.iter_enumerated() {
            let field = self.field(cx, field.as_usize());
            field.add_data_ranges(cx, offset, &mut this);
        }

        self.for_variant(cx, variant_index).add_data_ranges(cx, Size::ZERO, &mut this);

        // Padding specific to this variant: data in some variant, but not in this one.
        any.difference(&this).0.iter().map(|&(offset, size)| offset..offset + size).collect()
    }

    /// Extend `out` with all ranges of bytes that *may* carry relevant data for values of this type.
    /// For enums and unions there are offsets that are initialized for some
    /// variants but not for others; those offset *will* get added to `out`.
    fn add_data_ranges<C>(self, cx: &C, base_offset: Size, out: &mut RangeSet<Size>)
    where
        Ty: TyAbiInterface<'a, C> + Copy,
    {
        if self.is_zst() {
            return;
        }

        // Visit the fields of this value. For enum values the fields include the discriminant.
        match &self.fields {
            FieldsShape::Primitive => {
                out.add_range(base_offset, self.size);
            }
            &FieldsShape::Union(field_count) => {
                for field in 0..field_count.get() {
                    let field = self.field(cx, field);
                    field.add_data_ranges(cx, base_offset, out);
                }
            }
            &FieldsShape::Array { stride, count } => {
                let elem = self.field(cx, 0);

                // For scalars we know there is no padding between the elements,
                // so the entire array is a single big data range.
                if elem.backend_repr.is_scalar() {
                    out.add_range(base_offset, elem.size * count);
                } else {
                    // FIXME: this is really inefficient for large arrays.
                    for idx in 0..count {
                        elem.add_data_ranges(cx, base_offset + idx * stride, out);
                    }
                }
            }
            FieldsShape::Arbitrary { offsets, in_memory_order: _ } => {
                for (field, &offset) in offsets.iter_enumerated() {
                    let field = self.field(cx, field.as_usize());
                    field.add_data_ranges(cx, base_offset + offset, out);
                }
            }
        }

        // Visit the fields of each variant.
        match &self.variants {
            Variants::Empty | Variants::Single { index: _ } => { /* done */ }
            Variants::Multiple { variants, .. } => {
                for variant in variants.indices() {
                    let variant = self.for_variant(cx, variant);
                    variant.add_data_ranges(cx, base_offset, out);
                }
            }
        }
    }
}