//! RISC-V general purpose registers

#[cfg(test)]
mod tests;

use crate::registers::general_purpose::private::PhantomRegister;
use core::fmt;
use core::hint::unreachable_unchecked;
use core::marker::Destruct;
use core::ops::{
    Add, AddAssign, BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Not, Shl, Shr,
    Sub, SubAssign,
};

mod private {
    use core::marker::PhantomData;

    #[derive(Debug, Clone, Copy)]
    pub struct PhantomRegister<Type>(PhantomData<Type>);
}

/// Register type.
///
/// `u32` for RV32 and `u64` for RV64.
pub const trait RegType
where
    Self: [const] Default
        + [const] From<bool>
        + [const] From<u8>
        + [const] From<u16>
        + [const] From<u32>
        + [const] Eq
        + [const] Ord
        + [const] Add<Output = Self>
        + [const] AddAssign
        + [const] Sub<Output = Self>
        + [const] SubAssign
        + [const] BitAnd<Output = Self>
        + [const] BitAndAssign
        + [const] BitOr<Output = Self>
        + [const] BitOrAssign
        + [const] BitXor<Output = Self>
        + [const] BitXorAssign
        + [const] Not<Output = Self>
        + [const] Shl<u8, Output = Self>
        + [const] Shl<u16, Output = Self>
        + [const] Shl<u32, Output = Self>
        + [const] Shl<i32, Output = Self>
        + [const] Shr<u8, Output = Self>
        + [const] Shr<u16, Output = Self>
        + [const] Shr<u32, Output = Self>
        + [const] Shr<i32, Output = Self>
        + fmt::Display
        + fmt::LowerHex
        + fmt::UpperHex
        + fmt::Debug
        + Copy
        + Send
        + Sync
        + Sized
        + 'static,
{
    /// The size of this type in bits
    const BITS: u8;

    /// Convert to `u64`
    fn as_u64(&self) -> u64;
}

impl const RegType for u32 {
    const BITS: u8 = u32::BITS as u8;

    #[inline(always)]
    fn as_u64(&self) -> u64 {
        u64::from(*self)
    }
}

impl const RegType for u64 {
    const BITS: u8 = u64::BITS as u8;

    #[inline(always)]
    fn as_u64(&self) -> u64 {
        *self
    }
}

/// GPR (General Purpose Register)
///
/// # Safety
/// `Self::offset()` must return values in `0..Self::N` range. `Self::from_bits()` must return
/// `Some()` for `0..=31` if `Self::RVE = false` and `0..=15` if `Self::RVE = true`.
pub const unsafe trait Register:
    fmt::Display + fmt::Debug + [const] Eq + [const] Destruct + Copy + Send + Sync + Sized + 'static
{
    /// Whether this is RVE variant with the number of general purpose registers reduced to 16
    const RVE: bool;
    // TODO: Get rid of this field, replace with a separate trait that `Registers` will use to get
    //  the number of "slots" to store registers in
    /// The number of general purpose registers.
    ///
    /// Canonically 32 unless E extension is used, in which case 16.
    const N: usize;
    /// XLEN
    const XLEN: u8 = Self::Type::BITS;
    /// Zero register
    const ZERO: Self;
    /// Stack pointer register
    const SP: Self;
    /// Return address register
    const RA: Self;
    /// Function argument register a0
    const A0: Self;
    /// Function argument register a1
    const A1: Self;
    /// Register type.
    ///
    /// `u32` for RV32 and `u64` for RV64.
    type Type: [const] RegType;

    /// Create a register from its bit representation
    fn from_bits(bits: u8) -> Option<Self>;

    /// Offset in a set of registers
    fn offset(self) -> u8;
}

/// A set of RISC-V GPRs (General Purpose Registers)
#[derive(Debug, Clone, Copy)]
#[repr(align(16))]
pub struct Registers<Reg>
where
    Reg: Register,
    [(); Reg::N]:,
{
    regs: [Reg::Type; Reg::N],
}

impl<Reg> Default for Registers<Reg>
where
    Reg: Register,
    [(); Reg::N]:,
{
    #[inline(always)]
    fn default() -> Self {
        Self {
            regs: [Reg::Type::default(); Reg::N],
        }
    }
}

const impl<Reg> Registers<Reg>
where
    Reg: Register,
    [(); Reg::N]:,
{
    /// Read register value
    #[inline(always)]
    pub fn read(&self, reg: Reg) -> Reg::Type
    where
        Reg: [const] Register,
    {
        if reg == Reg::ZERO {
            // Always zero
            return Reg::Type::default();
        }

        // SAFETY: register offset is always within bounds
        *unsafe { self.regs.get_unchecked(usize::from(reg.offset())) }
    }

    /// Write register value
    #[inline(always)]
    pub fn write(&mut self, reg: Reg, value: Reg::Type)
    where
        Reg: [const] Register,
    {
        if reg == Reg::ZERO {
            // Writes are ignored
            return;
        }

        // SAFETY: register offset is always within bounds
        *unsafe { self.regs.get_unchecked_mut(usize::from(reg.offset())) } = value;
    }
}

/// RISC-V general purpose register for RV32E/RV64E.
///
/// Use `Type = u32` for RV32E and `Type = u64` for RV64E.
#[derive(Clone, Copy)]
#[repr(u8)]
pub enum EReg<Type> {
    /// Always zero: `x0`
    Zero = 0,
    /// Return address: `x1`
    Ra = 1,
    /// Stack pointer: `x2`
    Sp = 2,
    /// Global pointer: `x3`
    Gp = 3,
    /// Thread pointer: `x4`
    Tp = 4,
    /// Temporary/alternate return address: `x5`
    T0 = 5,
    /// Temporary: `x6`
    T1 = 6,
    /// Temporary: `x7`
    T2 = 7,
    /// Saved register/frame pointer: `x8`
    S0 = 8,
    /// Saved register: `x9`
    S1 = 9,
    /// Function argument/return value: `x10`
    A0 = 10,
    /// Function argument/return value: `x11`
    A1 = 11,
    /// Function argument: `x12`
    A2 = 12,
    /// Function argument: `x13`
    A3 = 13,
    /// Function argument: `x14`
    A4 = 14,
    /// Function argument: `x15`
    A5 = 15,
    /// Phantom register that is never constructed and is only used due to type system limitations
    #[doc(hidden)]
    Phantom(PhantomRegister<Type>),
}

impl<Type> fmt::Display for EReg<Type> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Zero => write!(f, "zero"),
            Self::Ra => write!(f, "ra"),
            Self::Sp => write!(f, "sp"),
            Self::Gp => write!(f, "gp"),
            Self::Tp => write!(f, "tp"),
            Self::T0 => write!(f, "t0"),
            Self::T1 => write!(f, "t1"),
            Self::T2 => write!(f, "t2"),
            Self::S0 => write!(f, "s0"),
            Self::S1 => write!(f, "s1"),
            Self::A0 => write!(f, "a0"),
            Self::A1 => write!(f, "a1"),
            Self::A2 => write!(f, "a2"),
            Self::A3 => write!(f, "a3"),
            Self::A4 => write!(f, "a4"),
            Self::A5 => write!(f, "a5"),
            Self::Phantom(_) => {
                // SAFETY: Phantom register is never constructed
                unsafe { unreachable_unchecked() }
            }
        }
    }
}

impl<Type> fmt::Debug for EReg<Type> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(self, f)
    }
}

impl<Type> const PartialEq for EReg<Type> {
    #[inline(always)]
    fn eq(&self, other: &Self) -> bool {
        // This is quite ugly, but there doesn't seem to be a much better way with `Phantom` variant
        matches!(
            (self, other),
            (Self::Zero, Self::Zero)
                | (Self::Ra, Self::Ra)
                | (Self::Sp, Self::Sp)
                | (Self::Gp, Self::Gp)
                | (Self::Tp, Self::Tp)
                | (Self::T0, Self::T0)
                | (Self::T1, Self::T1)
                | (Self::T2, Self::T2)
                | (Self::S0, Self::S0)
                | (Self::S1, Self::S1)
                | (Self::A0, Self::A0)
                | (Self::A1, Self::A1)
                | (Self::A2, Self::A2)
                | (Self::A3, Self::A3)
                | (Self::A4, Self::A4)
                | (Self::A5, Self::A5)
                | (Self::Phantom(_), Self::Phantom(_))
        )
    }
}

impl<Type> const Eq for EReg<Type> {}

// SAFETY: `Self::offset()` returns values within `0..Self::N` range and `Self::from_bits()` returns
// `Some()` for `0..=15`
unsafe impl const Register for EReg<u32> {
    const RVE: bool = true;
    const N: usize = 16;
    const ZERO: Self = Self::Zero;
    const SP: Self = Self::Sp;
    const RA: Self = Self::Ra;
    const A0: Self = Self::A0;
    const A1: Self = Self::A1;
    type Type = u32;

    #[inline(always)]
    fn from_bits(bits: u8) -> Option<Self> {
        match bits {
            0 => Some(Self::Zero),
            1 => Some(Self::Ra),
            2 => Some(Self::Sp),
            3 => Some(Self::Gp),
            4 => Some(Self::Tp),
            5 => Some(Self::T0),
            6 => Some(Self::T1),
            7 => Some(Self::T2),
            8 => Some(Self::S0),
            9 => Some(Self::S1),
            10 => Some(Self::A0),
            11 => Some(Self::A1),
            12 => Some(Self::A2),
            13 => Some(Self::A3),
            14 => Some(Self::A4),
            15 => Some(Self::A5),
            _ => None,
        }
    }

    #[inline(always)]
    fn offset(self) -> u8 {
        // SAFETY: Enum is `#[repr(u8)]` and doesn't have any fields
        unsafe { core::mem::transmute::<Self, u8>(self) }
    }
}

// SAFETY: `Self::offset()` returns values within `0..Self::N` range and `Self::from_bits()` returns
// `Some()` for `0..=15`
unsafe impl const Register for EReg<u64> {
    const RVE: bool = true;
    const N: usize = 16;
    const ZERO: Self = Self::Zero;
    const SP: Self = Self::Sp;
    const RA: Self = Self::Ra;
    const A0: Self = Self::A0;
    const A1: Self = Self::A1;
    type Type = u64;

    #[inline(always)]
    fn from_bits(bits: u8) -> Option<Self> {
        match bits {
            0 => Some(Self::Zero),
            1 => Some(Self::Ra),
            2 => Some(Self::Sp),
            3 => Some(Self::Gp),
            4 => Some(Self::Tp),
            5 => Some(Self::T0),
            6 => Some(Self::T1),
            7 => Some(Self::T2),
            8 => Some(Self::S0),
            9 => Some(Self::S1),
            10 => Some(Self::A0),
            11 => Some(Self::A1),
            12 => Some(Self::A2),
            13 => Some(Self::A3),
            14 => Some(Self::A4),
            15 => Some(Self::A5),
            _ => None,
        }
    }

    #[inline(always)]
    fn offset(self) -> u8 {
        // SAFETY: Enum is `#[repr(u8)]` and doesn't have any fields
        unsafe { core::mem::transmute::<Self, u8>(self) }
    }
}

// TODO: Reorder `t*` and `s*` registers to be next to each other for performance reasons?
/// RISC-V general purpose register for RV32I/RV64I.
///
/// Use `Type = u32` for RV32I and `Type = u64` for RV64I.
#[derive(Clone, Copy)]
#[repr(u8)]
pub enum Reg<Type> {
    /// Always zero: `x0`
    Zero = 0,
    /// Return address: `x1`
    Ra = 1,
    /// Stack pointer: `x2`
    Sp = 2,
    /// Global pointer: `x3`
    Gp = 3,
    /// Thread pointer: `x4`
    Tp = 4,
    /// Temporary/alternate return address: `x5`
    T0 = 5,
    /// Temporary: `x6`
    T1 = 6,
    /// Temporary: `x7`
    T2 = 7,
    /// Saved register/frame pointer: `x8`
    S0 = 8,
    /// Saved register: `x9`
    S1 = 9,
    /// Function argument/return value: `x10`
    A0 = 10,
    /// Function argument/return value: `x11`
    A1 = 11,
    /// Function argument: `x12`
    A2 = 12,
    /// Function argument: `x13`
    A3 = 13,
    /// Function argument: `x14`
    A4 = 14,
    /// Function argument: `x15`
    A5 = 15,
    /// Function argument: `x16`
    A6 = 16,
    /// Function argument: `x17`
    A7 = 17,
    /// Saved register: `x18`
    S2 = 18,
    /// Saved register: `x19`
    S3 = 19,
    /// Saved register: `x20`
    S4 = 20,
    /// Saved register: `x21`
    S5 = 21,
    /// Saved register: `x22`
    S6 = 22,
    /// Saved register: `x23`
    S7 = 23,
    /// Saved register: `x24`
    S8 = 24,
    /// Saved register: `x25`
    S9 = 25,
    /// Saved register: `x26`
    S10 = 26,
    /// Saved register: `x27`
    S11 = 27,
    /// Temporary: `x28`
    T3 = 28,
    /// Temporary: `x29`
    T4 = 29,
    /// Temporary: `x30`
    T5 = 30,
    /// Temporary: `x31`
    T6 = 31,
    /// Phantom register that is never constructed and is only used due to type system limitations
    #[doc(hidden)]
    Phantom(PhantomRegister<Type>),
}

impl<Type> const From<EReg<u64>> for Reg<Type> {
    #[inline(always)]
    fn from(reg: EReg<u64>) -> Self {
        match reg {
            EReg::Zero => Self::Zero,
            EReg::Ra => Self::Ra,
            EReg::Sp => Self::Sp,
            EReg::Gp => Self::Gp,
            EReg::Tp => Self::Tp,
            EReg::T0 => Self::T0,
            EReg::T1 => Self::T1,
            EReg::T2 => Self::T2,
            EReg::S0 => Self::S0,
            EReg::S1 => Self::S1,
            EReg::A0 => Self::A0,
            EReg::A1 => Self::A1,
            EReg::A2 => Self::A2,
            EReg::A3 => Self::A3,
            EReg::A4 => Self::A4,
            EReg::A5 => Self::A5,
            EReg::Phantom(_) => {
                // SAFETY: Phantom register is never constructed
                unsafe { unreachable_unchecked() }
            }
        }
    }
}

impl<Type> fmt::Display for Reg<Type> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Zero => write!(f, "zero"),
            Self::Ra => write!(f, "ra"),
            Self::Sp => write!(f, "sp"),
            Self::Gp => write!(f, "gp"),
            Self::Tp => write!(f, "tp"),
            Self::T0 => write!(f, "t0"),
            Self::T1 => write!(f, "t1"),
            Self::T2 => write!(f, "t2"),
            Self::S0 => write!(f, "s0"),
            Self::S1 => write!(f, "s1"),
            Self::A0 => write!(f, "a0"),
            Self::A1 => write!(f, "a1"),
            Self::A2 => write!(f, "a2"),
            Self::A3 => write!(f, "a3"),
            Self::A4 => write!(f, "a4"),
            Self::A5 => write!(f, "a5"),
            Self::A6 => write!(f, "a6"),
            Self::A7 => write!(f, "a7"),
            Self::S2 => write!(f, "s2"),
            Self::S3 => write!(f, "s3"),
            Self::S4 => write!(f, "s4"),
            Self::S5 => write!(f, "s5"),
            Self::S6 => write!(f, "s6"),
            Self::S7 => write!(f, "s7"),
            Self::S8 => write!(f, "s8"),
            Self::S9 => write!(f, "s9"),
            Self::S10 => write!(f, "s10"),
            Self::S11 => write!(f, "s11"),
            Self::T3 => write!(f, "t3"),
            Self::T4 => write!(f, "t4"),
            Self::T5 => write!(f, "t5"),
            Self::T6 => write!(f, "t6"),
            Self::Phantom(_) => {
                // SAFETY: Phantom register is never constructed
                unsafe { unreachable_unchecked() }
            }
        }
    }
}

impl<Type> fmt::Debug for Reg<Type> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(self, f)
    }
}

impl<Type> const PartialEq for Reg<Type> {
    #[inline(always)]
    fn eq(&self, other: &Self) -> bool {
        // This is quite ugly, but there doesn't seem to be a much better way with `Phantom` variant
        matches!(
            (self, other),
            (Self::Zero, Self::Zero)
                | (Self::Ra, Self::Ra)
                | (Self::Sp, Self::Sp)
                | (Self::Gp, Self::Gp)
                | (Self::Tp, Self::Tp)
                | (Self::T0, Self::T0)
                | (Self::T1, Self::T1)
                | (Self::T2, Self::T2)
                | (Self::S0, Self::S0)
                | (Self::S1, Self::S1)
                | (Self::A0, Self::A0)
                | (Self::A1, Self::A1)
                | (Self::A2, Self::A2)
                | (Self::A3, Self::A3)
                | (Self::A4, Self::A4)
                | (Self::A5, Self::A5)
                | (Self::A6, Self::A6)
                | (Self::A7, Self::A7)
                | (Self::S2, Self::S2)
                | (Self::S3, Self::S3)
                | (Self::S4, Self::S4)
                | (Self::S5, Self::S5)
                | (Self::S6, Self::S6)
                | (Self::S7, Self::S7)
                | (Self::S8, Self::S8)
                | (Self::S9, Self::S9)
                | (Self::S10, Self::S10)
                | (Self::S11, Self::S11)
                | (Self::T3, Self::T3)
                | (Self::T4, Self::T4)
                | (Self::T5, Self::T5)
                | (Self::T6, Self::T6)
                | (Self::Phantom(_), Self::Phantom(_))
        )
    }
}

impl<Type> const Eq for Reg<Type> {}

// SAFETY: `Self::offset()` returns values within `0..Self::N` range and `Self::from_bits()` returns
// `Some()` for `0..=31`
unsafe impl const Register for Reg<u32> {
    const RVE: bool = false;
    const N: usize = 32;
    const ZERO: Self = Self::Zero;
    const SP: Self = Self::Sp;
    const RA: Self = Self::Ra;
    const A0: Self = Self::A0;
    const A1: Self = Self::A1;
    type Type = u32;

    #[inline(always)]
    fn from_bits(bits: u8) -> Option<Self> {
        match bits {
            0 => Some(Self::Zero),
            1 => Some(Self::Ra),
            2 => Some(Self::Sp),
            3 => Some(Self::Gp),
            4 => Some(Self::Tp),
            5 => Some(Self::T0),
            6 => Some(Self::T1),
            7 => Some(Self::T2),
            8 => Some(Self::S0),
            9 => Some(Self::S1),
            10 => Some(Self::A0),
            11 => Some(Self::A1),
            12 => Some(Self::A2),
            13 => Some(Self::A3),
            14 => Some(Self::A4),
            15 => Some(Self::A5),
            16 => Some(Self::A6),
            17 => Some(Self::A7),
            18 => Some(Self::S2),
            19 => Some(Self::S3),
            20 => Some(Self::S4),
            21 => Some(Self::S5),
            22 => Some(Self::S6),
            23 => Some(Self::S7),
            24 => Some(Self::S8),
            25 => Some(Self::S9),
            26 => Some(Self::S10),
            27 => Some(Self::S11),
            28 => Some(Self::T3),
            29 => Some(Self::T4),
            30 => Some(Self::T5),
            31 => Some(Self::T6),
            _ => None,
        }
    }

    #[inline(always)]
    fn offset(self) -> u8 {
        // SAFETY: Enum is `#[repr(u8)]` and doesn't have any fields
        unsafe { core::mem::transmute::<Self, u8>(self) }
    }
}

// SAFETY: `Self::offset()` returns values within `0..Self::N` range and `Self::from_bits()` returns
// `Some()` for `0..=31`
unsafe impl const Register for Reg<u64> {
    const RVE: bool = false;
    const N: usize = 32;
    const ZERO: Self = Self::Zero;
    const SP: Self = Self::Sp;
    const RA: Self = Self::Ra;
    const A0: Self = Self::A0;
    const A1: Self = Self::A1;
    type Type = u64;

    #[inline(always)]
    fn from_bits(bits: u8) -> Option<Self> {
        match bits {
            0 => Some(Self::Zero),
            1 => Some(Self::Ra),
            2 => Some(Self::Sp),
            3 => Some(Self::Gp),
            4 => Some(Self::Tp),
            5 => Some(Self::T0),
            6 => Some(Self::T1),
            7 => Some(Self::T2),
            8 => Some(Self::S0),
            9 => Some(Self::S1),
            10 => Some(Self::A0),
            11 => Some(Self::A1),
            12 => Some(Self::A2),
            13 => Some(Self::A3),
            14 => Some(Self::A4),
            15 => Some(Self::A5),
            16 => Some(Self::A6),
            17 => Some(Self::A7),
            18 => Some(Self::S2),
            19 => Some(Self::S3),
            20 => Some(Self::S4),
            21 => Some(Self::S5),
            22 => Some(Self::S6),
            23 => Some(Self::S7),
            24 => Some(Self::S8),
            25 => Some(Self::S9),
            26 => Some(Self::S10),
            27 => Some(Self::S11),
            28 => Some(Self::T3),
            29 => Some(Self::T4),
            30 => Some(Self::T5),
            31 => Some(Self::T6),
            _ => None,
        }
    }

    #[inline(always)]
    fn offset(self) -> u8 {
        // SAFETY: Enum is `#[repr(u8)]` and doesn't have any fields
        unsafe { core::mem::transmute::<Self, u8>(self) }
    }
}
