//! The FPU, SSE and AVX registers of user programs (stage 15.1).
//!
//! Besides the 16 general-purpose registers, an x86_64 CPU has a second set
//! for floating point and vector maths: the old x87 FPU stack, the 16 SSE
//! registers xmm0-xmm15 (128 bits each) plus their control register MXCSR,
//! and, on CPUs with AVX, the upper halves that make them 256-bit ymm0-ymm15.
//! Every C compiler uses xmm registers for `double` arithmetic and for fast
//! `memcpy`, so to the kernel they are as much "the program's registers" as
//! rax is: if two threads share a CPU, each must find its own values there
//! when it runs again.
//!
//! The kernel itself is built *soft-float* (see x86_64-rust_kernel.json): it
//! never touches these registers. That makes the job simple: whatever is in
//! them while the kernel runs belongs to the user thread the CPU last ran.
//! So we only
//!
//!   * save them into the thread's `FpuState` when the scheduler switches
//!     away from a user thread, and load the next user thread's (thread.rs,
//!     "eager" switching -- simple, and cheap enough at ~1 KiB);
//!   * copy them into a fork's child / a new thread (they start with the
//!     parent's values, like Linux);
//!   * reset them on execve (a new program starts clean);
//!   * put them in signal frames, so a handler may use floats and the
//!     interrupted code gets its own values back (signal.rs).
//!
//! The CPU has instructions that save/load all of it at once:
//!
//!   * FXSAVE / FXRSTOR (every x86_64 CPU): x87 + SSE, a 512-byte area.
//!   * XSAVE / XRSTOR (newer CPUs): a 512-byte legacy part in the same
//!     format, a 64-byte header (which parts are saved), then one part per
//!     extra feature -- AVX's upper ymm halves are 256 bytes more. Which
//!     features are switched on is the OS's choice, in register XCR0.
//!
//! Before user code may use any of this, the OS must say it will save the
//! state (otherwise the CPU raises "invalid opcode"): control register bits
//! CR4.OSFXSR (SSE allowed), CR4.OSXMMEXCPT (SSE maths errors raise #XM, not
//! #UD) and CR4.OSXSAVE (XSAVE + XCR0 allowed), on every CPU.

use core::arch::asm;
use core::arch::x86_64::{__cpuid, __cpuid_count};
use core::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};

/// Room for the biggest state we enable: x87 + SSE + AVX = 832 bytes with
/// XSAVE (we don't switch on AVX-512, which would need ~2.7 KiB).
pub const AREA_SIZE: usize = 1024;

/// One thread's saved FPU/SSE/AVX registers, in the FXSAVE/XSAVE memory
/// format. XSAVE needs 64-byte alignment.
#[repr(C, align(64))]
#[derive(Clone)]
pub struct FpuState(pub [u8; AREA_SIZE]);

/// XSAVE is used (else FXSAVE).
static USE_XSAVE: AtomicBool = AtomicBool::new(false);
/// The features in XCR0 (bit 0 x87, bit 1 SSE, bit 2 AVX).
static XCR0: AtomicU64 = AtomicU64::new(0);
/// Bytes of state: 512 (FXSAVE) or what CPUID says for our XCR0.
static SIZE: AtomicU32 = AtomicU32::new(512);
/// MXCSR bits the CPU supports; loading any other bit faults (#GP).
static MXCSR_MASK: AtomicU32 = AtomicU32::new(0xffbf);

/// MXCSR at reset: all six SSE exceptions masked, round to nearest.
pub const MXCSR_DEFAULT: u32 = 0x1f80;
/// x87 control word at reset (FNINIT): all exceptions masked, 64-bit precision.
const FCW_DEFAULT: u16 = 0x037f;

pub const XFEATURE_X87: u64 = 1;
pub const XFEATURE_SSE: u64 = 2;
pub const XFEATURE_AVX: u64 = 4;

/// Switch the FPU/SSE (and AVX) on for this CPU. Every CPU calls it once at
/// boot: CR0/CR4/XCR0 are per CPU.
pub fn init_cpu() {
    // CPUID leaf 1: ECX bit 26 = XSAVE, bit 28 = AVX.
    let leaf1 = __cpuid(1);
    let has_xsave = leaf1.ecx & (1 << 26) != 0;
    let has_avx = leaf1.ecx & (1 << 28) != 0;
    unsafe {
        // CR0: EM (bit 2, "no FPU: trap every FPU instruction") and TS (bit
        // 3, "trap the next one", for lazy switching) off; MP (bit 1) and NE
        // (bit 5, report x87 errors as exception #MF, not via the old PC IRQ 13) on.
        let mut cr0: u64;
        asm!("mov {}, cr0", out(reg) cr0);
        cr0 = (cr0 & !(1 << 2 | 1 << 3)) | 1 << 1 | 1 << 5;
        asm!("mov cr0, {}", in(reg) cr0);
        // CR4: OSFXSR (bit 9) + OSXMMEXCPT (bit 10), and OSXSAVE (bit 18).
        let mut cr4: u64;
        asm!("mov {}, cr4", out(reg) cr4);
        cr4 |= 1 << 9 | 1 << 10;
        if has_xsave {
            cr4 |= 1 << 18;
        }
        asm!("mov cr4, {}", in(reg) cr4);
    }
    let xcr0 = XFEATURE_X87 | XFEATURE_SSE | if has_avx && has_xsave { XFEATURE_AVX } else { 0 };
    if has_xsave {
        // XSETBV ecx=0: write XCR0 (edx:eax).
        unsafe { asm!("xsetbv", in("ecx") 0, in("eax") xcr0 as u32, in("edx") (xcr0 >> 32) as u32) };
        // CPUID leaf 0xD, sub-leaf 0: EBX = bytes XSAVE needs for the
        // features currently on in XCR0.
        let size = __cpuid_count(0xd, 0).ebx;
        assert!(size as usize <= AREA_SIZE, "XSAVE area too big: {}", size);
        SIZE.store(size, Ordering::Relaxed);
        XCR0.store(xcr0, Ordering::Relaxed);
        USE_XSAVE.store(true, Ordering::Relaxed);
    } else {
        XCR0.store(XFEATURE_X87 | XFEATURE_SSE, Ordering::Relaxed);
    }
    // Clean registers, and find out which MXCSR bits are allowed: FXSAVE
    // stores that mask at byte 28 (0 means the old default, 0xffbf).
    unsafe { asm!("fninit") };
    let mut probe = FpuState([0; AREA_SIZE]);
    unsafe { asm!("fxsave64 [{}]", in(reg) probe.0.as_mut_ptr()) };
    let mask = u32::from_le_bytes(probe.0[28..32].try_into().unwrap());
    MXCSR_MASK.store(if mask == 0 { 0xffbf } else { mask }, Ordering::Relaxed);
    restore(&initial());
}

pub fn uses_xsave() -> bool {
    USE_XSAVE.load(Ordering::Relaxed)
}
pub fn xfeatures() -> u64 {
    XCR0.load(Ordering::Relaxed)
}
/// One line for the boot log, e.g. "XSAVE: x87 SSE AVX, 832 bytes per thread".
pub fn describe() -> alloc::string::String {
    let avx = if xfeatures() & XFEATURE_AVX != 0 { " AVX" } else { "" };
    let how = if uses_xsave() { "XSAVE" } else { "FXSAVE" };
    alloc::format!("{}: x87 SSE{}, {} bytes per thread", how, avx, size())
}

/// Bytes of saved state (512 with FXSAVE).
pub fn size() -> usize {
    SIZE.load(Ordering::Relaxed) as usize
}

/// The state a new program starts with: everything zero, the x87 control
/// word and MXCSR at their reset values. The XSAVE header (bytes 512..576)
/// is all zero, which tells XRSTOR "every extra feature is in its initial
/// state" (zero).
pub fn initial() -> FpuState {
    let mut s = FpuState([0; AREA_SIZE]);
    s.0[0..2].copy_from_slice(&FCW_DEFAULT.to_le_bytes());
    s.0[24..28].copy_from_slice(&MXCSR_DEFAULT.to_le_bytes());
    if uses_xsave() {
        // ... except x87 + SSE, which we say are valid (the values above).
        s.0[512..520].copy_from_slice(&(XFEATURE_X87 | XFEATURE_SSE).to_le_bytes());
    }
    s
}

/// Save this CPU's registers into `s`.
pub fn save(s: &mut FpuState) {
    let p = s.0.as_mut_ptr();
    unsafe {
        if uses_xsave() {
            let m = xfeatures();
            asm!("xsave64 [{}]", in(reg) p, in("eax") m as u32, in("edx") (m >> 32) as u32);
        } else {
            asm!("fxsave64 [{}]", in(reg) p);
        }
    }
}

/// The registers right now (for fork, clone and signal frames).
pub fn save_live() -> FpuState {
    let mut s = FpuState([0; AREA_SIZE]);
    save(&mut s);
    s
}

/// Load `s` into this CPU's registers. `s` must be valid (see `sanitize`):
/// a bad MXCSR or XSAVE header would fault in the kernel.
pub fn restore(s: &FpuState) {
    let p = s.0.as_ptr();
    unsafe {
        if uses_xsave() {
            let m = xfeatures();
            asm!("xrstor64 [{}]", in(reg) p, in("eax") m as u32, in("edx") (m >> 32) as u32);
        } else {
            asm!("fxrstor64 [{}]", in(reg) p);
        }
    }
}

/// execve and signal handlers start from clean registers.
pub fn reset_live() {
    restore(&initial());
}

/// Make state that came from user memory (a signal frame the program may
/// have edited) safe to load: false if MXCSR has bits the CPU doesn't
/// support (XRSTOR would fault). Fixes the XSAVE header: only features we
/// enabled, standard (not compacted) format, reserved bytes zero.
pub fn sanitize(s: &mut FpuState) -> bool {
    let mxcsr = u32::from_le_bytes(s.0[24..28].try_into().unwrap());
    if mxcsr & !MXCSR_MASK.load(Ordering::Relaxed) != 0 {
        return false;
    }
    if uses_xsave() {
        let bv = u64::from_le_bytes(s.0[512..520].try_into().unwrap()) & xfeatures();
        s.0[512..520].copy_from_slice(&bv.to_le_bytes());
        s.0[520..576].fill(0); // XCOMP_BV + reserved
    }
    true
}

/// The SSE status flags of the last SIMD exception (#XM), for SIGFPE's
/// si_code: MXCSR bits 0-5 = invalid, denormal, divide by zero, overflow,
/// underflow, precision.
pub fn mxcsr() -> u32 {
    let mut v: u32 = 0;
    unsafe { asm!("stmxcsr [{}]", in(reg) &mut v as *mut u32) };
    v
}