//! Compact semantic types for Slonik IR.

use core::fmt;

/// A semantic type in Slonik.
///
/// `Type` describes the meaning and shape of an SSA value.
/// Integers are signedness-agnostic; signedness is carried by the operation
/// (e.g. `sdiv` vs `udiv`, `icmp slt` vs `icmp ult`).
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum Type {
    /// The absence of a value.
    ///
    /// This should generally not appear as the type of an SSA `Value`.
    Void,

    /// A one-bit boolean.
    Bool,

    /// A signedness-agnostic integer of the given bit width.
    I(u16),

    /// A floating-point scalar.
    Float(FloatType),

    /// A pointer value.
    Ptr(PtrType),

    /// A fixed-length or scalable SIMD vector.
    Vector(VectorType),
}

/// A floating-point scalar type.
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum FloatType {
    F16,
    F32,
    F64,
    F80,
    F128,
}

impl FloatType {
    /// Returns the width of this floating-point type in bits.
    pub const fn bits(self) -> u16 {
        match self {
            Self::F16 => 16,
            Self::F32 => 32,
            Self::F64 => 64,
            Self::F80 => 80,
            Self::F128 => 128,
        }
    }
}

/// A pointer type.
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct PtrType {
    bits: u16,
    addr_space: u16,
}

impl PtrType {
    /// Creates a new pointer type in address space 0.
    ///
    /// # Panics
    ///
    /// Panics if `bits == 0`.
    pub const fn new(bits: u16) -> Self {
        Self::with_addr_space(bits, 0)
    }

    /// Creates a new pointer type with an explicit address space.
    ///
    /// # Panics
    ///
    /// Panics if `bits == 0`.
    pub const fn with_addr_space(bits: u16, addr_space: u16) -> Self {
        assert!(bits != 0, "pointer type must have nonzero width");
        Self { bits, addr_space }
    }

    /// Returns the pointer width in bits.
    pub const fn bits(self) -> u16 {
        self.bits
    }

    /// Returns the address space.
    pub const fn addr_space(self) -> u16 {
        self.addr_space
    }
}

/// A scalar lane type for a vector.
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum LaneType {
    Bool,
    I(u16),
    Float(FloatType),
    Ptr(PtrType),
}

impl LaneType {
    /// Returns the width of one lane in bits.
    pub const fn bits(self) -> u16 {
        match self {
            Self::Bool => 1,
            Self::I(bits) => bits,
            Self::Float(ty) => ty.bits(),
            Self::Ptr(ty) => ty.bits(),
        }
    }

    /// Converts this lane type into the corresponding scalar `Type`.
    pub const fn as_type(self) -> Type {
        match self {
            Self::Bool => Type::Bool,
            Self::I(bits) => Type::I(bits),
            Self::Float(ty) => Type::Float(ty),
            Self::Ptr(ty) => Type::Ptr(ty),
        }
    }
}

/// A SIMD vector type.
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct VectorType {
    lanes: u16,
    lane: LaneType,
    scalable: bool,
}

impl VectorType {
    /// Creates a new fixed-length vector type.
    ///
    /// # Panics
    ///
    /// Panics if `lanes == 0`.
    pub const fn new(lane: LaneType, lanes: u16) -> Self {
        assert!(lanes != 0, "vector type must have at least one lane");
        Self {
            lanes,
            lane,
            scalable: false,
        }
    }

    /// Returns the number of lanes.
    pub const fn lanes(self) -> u16 {
        self.lanes
    }

    /// Returns the lane type.
    pub const fn lane(self) -> LaneType {
        self.lane
    }

    /// Returns whether this is a scalable vector.
    pub const fn scalable(self) -> bool {
        self.scalable
    }

    /// Returns the width of one lane in bits.
    pub const fn lane_bits(self) -> u16 {
        self.lane.bits()
    }

    /// Returns the minimum number of bits in the vector.
    pub const fn min_bits(self) -> u32 {
        (self.lanes as u32) * (self.lane.bits() as u32)
    }
}

impl Type {
    /// Returns `void`.
    pub const fn void() -> Self {
        Self::Void
    }

    /// Returns `bool`.
    pub const fn bool() -> Self {
        Self::Bool
    }

    /// Returns an integer type with the given width.
    pub const fn int(bits: u16) -> Self {
        Self::I(bits)
    }

    /// Returns a floating-point type.
    pub const fn float(kind: FloatType) -> Self {
        Self::Float(kind)
    }

    /// Returns a pointer type in address space 0.
    pub const fn ptr(bits: u16) -> Self {
        Self::Ptr(PtrType::new(bits))
    }

    /// Returns a pointer type with an explicit address space.
    pub const fn ptr_in(bits: u16, addr_space: u16) -> Self {
        Self::Ptr(PtrType::with_addr_space(bits, addr_space))
    }

    /// Returns a fixed-length vector type.
    pub const fn vector(lane: LaneType, lanes: u16) -> Self {
        Self::Vector(VectorType::new(lane, lanes))
    }

    /// Returns whether this is `void`.
    pub const fn is_void(self) -> bool {
        matches!(self, Self::Void)
    }

    /// Returns whether this is `bool`.
    pub const fn is_bool(self) -> bool {
        matches!(self, Self::Bool)
    }

    /// Returns whether this is an integer.
    pub const fn is_int(self) -> bool {
        matches!(self, Self::I(_))
    }

    /// Returns whether this is a float.
    pub const fn is_float(self) -> bool {
        matches!(self, Self::Float(_))
    }

    /// Returns whether this is a pointer.
    pub const fn is_ptr(self) -> bool {
        matches!(self, Self::Ptr(_))
    }

    /// Returns whether this is a vector.
    pub const fn is_vector(self) -> bool {
        matches!(self, Self::Vector(_))
    }

    /// Returns whether this is any scalar type.
    pub const fn is_scalar(self) -> bool {
        !matches!(self, Self::Vector(_))
    }

    /// Returns the scalar width in bits, if this is a scalar type.
    pub const fn bit_size(self) -> Option<u16> {
        match self {
            Self::Void => None,
            Self::Bool => Some(1),
            Self::I(bits) => Some(bits),
            Self::Float(ty) => Some(ty.bits()),
            Self::Ptr(ty) => Some(ty.bits()),
            Self::Vector(_) => None,
        }
    }

    /// Returns the vector description, if this is a vector.
    pub const fn as_vector(self) -> Option<VectorType> {
        match self {
            Self::Vector(ty) => Some(ty),
            _ => None,
        }
    }

    /// Returns the number of lanes if this is a vector.
    pub const fn lanes(self) -> Option<u16> {
        match self {
            Self::Vector(ty) => Some(ty.lanes()),
            _ => None,
        }
    }

    /// Convenience constructor for `i8`.
    pub const fn i8() -> Self {
        Self::I(8)
    }

    /// Convenience constructor for `i16`.
    pub const fn i16() -> Self {
        Self::I(16)
    }

    /// Convenience constructor for `i32`.
    pub const fn i32() -> Self {
        Self::I(32)
    }

    /// Convenience constructor for `i64`.
    pub const fn i64() -> Self {
        Self::I(64)
    }

    /// Convenience constructor for `i128`.
    pub const fn i128() -> Self {
        Self::I(128)
    }

    /// Convenience constructor for `f32`.
    pub const fn f32() -> Self {
        Self::float(FloatType::F32)
    }

    /// Convenience constructor for `f64`.
    pub const fn f64() -> Self {
        Self::float(FloatType::F64)
    }
}

impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Type::Void => f.write_str("void"),
            Type::Bool => f.write_str("bool"),
            Type::I(bits) => write!(f, "i{bits}"),
            Type::Float(FloatType::F16) => f.write_str("f16"),
            Type::Float(FloatType::F32) => f.write_str("f32"),
            Type::Float(FloatType::F64) => f.write_str("f64"),
            Type::Float(FloatType::F80) => f.write_str("f80"),
            Type::Float(FloatType::F128) => f.write_str("f128"),
            Type::Ptr(ty) => {
                if ty.addr_space() == 0 {
                    write!(f, "ptr{}", ty.bits())
                } else {
                    write!(f, "ptr{}@{}", ty.bits(), ty.addr_space())
                }
            }
            Type::Vector(ty) => {
                if ty.scalable() {
                    write!(f, "vec<?x{} x {}>", ty.lanes(), ty.lane())
                } else {
                    write!(f, "vec<{} x {}>", ty.lanes(), ty.lane())
                }
            }
        }
    }
}

impl fmt::Display for LaneType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            LaneType::Bool => f.write_str("bool"),
            LaneType::I(bits) => write!(f, "i{bits}"),
            LaneType::Float(FloatType::F16) => f.write_str("f16"),
            LaneType::Float(FloatType::F32) => f.write_str("f32"),
            LaneType::Float(FloatType::F64) => f.write_str("f64"),
            LaneType::Float(FloatType::F80) => f.write_str("f80"),
            LaneType::Float(FloatType::F128) => f.write_str("f128"),
            LaneType::Ptr(ty) => {
                if ty.addr_space() == 0 {
                    write!(f, "ptr{}", ty.bits())
                } else {
                    write!(f, "ptr{}@{}", ty.bits(), ty.addr_space())
                }
            }
        }
    }
}

pub const I32: Type = Type::i32();
