//! fputest -- every thread keeps its own FPU/SSE/AVX registers (stage 15.1).
//!
//! The kernel never uses the FPU registers itself, but user code does (any
//! C program's `double` maths, and libc's memcpy/strlen use SSE). So on a
//! context switch the kernel saves one thread's registers and loads the
//! next one's (XSAVE, or FXSAVE on older CPUs), a signal handler gets
//! clean registers and sigreturn puts the old ones back, fork copies them
//! and exec clears them. This program checks all of that:
//!
//! - threads + processes (spread over all CPUs) each load *distinct*
//!   values into xmm0-15 (all 256 bits of ymm0-15 if the CPU has AVX) and
//!   MXCSR, spin and yield for a while, then check nothing changed;
//! - a signal handler that scribbles over every register, synchronously
//!   (kill to ourselves) and asynchronously (another thread's tgkill while
//!   we spin): the interrupted code must see its own values afterwards;
//! - the handler starts with clean registers, finds the saved ones via
//!   uc_mcontext.fpstate (Linux's layout), and a value it changes there is
//!   what sigreturn loads;
//! - an unmasked SSE divide-by-zero raises SIGFPE (FPE_FLTDIV);
//! - fork copies the registers; exec starts with zeros and MXCSR 0x1f80.
//!
//! Our programs are built soft-float, so Rust code here never touches the
//! vector registers: only the `asm!` blocks below do. (That's also why the
//! asm can't name xmm registers as operands: it uses memory instead.)
#![no_std]
#![no_main]

#[macro_use]
#[path = "ulib.rs"]
mod ulib;

extern crate alloc;

use alloc::sync::Arc;
use alloc::vec::Vec;
use core::arch::asm;
use core::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
use ulib::*;

/// The 16 vector registers, 32 bytes each (only the low 16 used without AVX).
#[repr(C, align(32))]
#[derive(Clone, Copy, PartialEq)]
struct Regs([u8; 512]);

static AVX: AtomicBool = AtomicBool::new(false);

/// Can we use AVX? The CPU must have it (CPUID.1:ECX bit 28) *and* the
/// kernel must have turned on XSAVE (OSXSAVE, bit 27) with the AVX state
/// in XCR0 (bits 1 and 2): otherwise AVX instructions fault (#UD).
fn detect_avx() -> bool {
    let ecx = core::arch::x86_64::__cpuid(1).ecx;
    if ecx & (1 << 27) == 0 || ecx & (1 << 28) == 0 {
        return false;
    }
    let (lo, _hi): (u32, u32);
    unsafe { asm!("xgetbv", in("ecx") 0, out("eax") lo, out("edx") _hi, options(nomem, nostack)) };
    lo & 6 == 6
}
fn avx() -> bool {
    AVX.load(Ordering::Relaxed)
}

/// Load all vector registers from `r`.
fn load(r: &Regs) {
    unsafe {
        if avx() {
            asm!(
            "vmovdqu ymm0, [{p} + 0]",
            "vmovdqu ymm1, [{p} + 32]",
            "vmovdqu ymm2, [{p} + 64]",
            "vmovdqu ymm3, [{p} + 96]",
            "vmovdqu ymm4, [{p} + 128]",
            "vmovdqu ymm5, [{p} + 160]",
            "vmovdqu ymm6, [{p} + 192]",
            "vmovdqu ymm7, [{p} + 224]",
            "vmovdqu ymm8, [{p} + 256]",
            "vmovdqu ymm9, [{p} + 288]",
            "vmovdqu ymm10, [{p} + 320]",
            "vmovdqu ymm11, [{p} + 352]",
            "vmovdqu ymm12, [{p} + 384]",
            "vmovdqu ymm13, [{p} + 416]",
            "vmovdqu ymm14, [{p} + 448]",
            "vmovdqu ymm15, [{p} + 480]",
                p = in(reg) r.0.as_ptr(), options(nostack, readonly, preserves_flags));
        } else {
            asm!(
            "movdqu xmm0, [{p} + 0]",
            "movdqu xmm1, [{p} + 32]",
            "movdqu xmm2, [{p} + 64]",
            "movdqu xmm3, [{p} + 96]",
            "movdqu xmm4, [{p} + 128]",
            "movdqu xmm5, [{p} + 160]",
            "movdqu xmm6, [{p} + 192]",
            "movdqu xmm7, [{p} + 224]",
            "movdqu xmm8, [{p} + 256]",
            "movdqu xmm9, [{p} + 288]",
            "movdqu xmm10, [{p} + 320]",
            "movdqu xmm11, [{p} + 352]",
            "movdqu xmm12, [{p} + 384]",
            "movdqu xmm13, [{p} + 416]",
            "movdqu xmm14, [{p} + 448]",
            "movdqu xmm15, [{p} + 480]",
                p = in(reg) r.0.as_ptr(), options(nostack, readonly, preserves_flags));
        }
    }
}

/// Store all vector registers (the bytes AVX doesn't have stay 0).
fn store() -> Regs {
    let mut r = Regs([0; 512]);
    unsafe {
        if avx() {
            asm!(
            "vmovdqu [{p} + 0], ymm0",
            "vmovdqu [{p} + 32], ymm1",
            "vmovdqu [{p} + 64], ymm2",
            "vmovdqu [{p} + 96], ymm3",
            "vmovdqu [{p} + 128], ymm4",
            "vmovdqu [{p} + 160], ymm5",
            "vmovdqu [{p} + 192], ymm6",
            "vmovdqu [{p} + 224], ymm7",
            "vmovdqu [{p} + 256], ymm8",
            "vmovdqu [{p} + 288], ymm9",
            "vmovdqu [{p} + 320], ymm10",
            "vmovdqu [{p} + 352], ymm11",
            "vmovdqu [{p} + 384], ymm12",
            "vmovdqu [{p} + 416], ymm13",
            "vmovdqu [{p} + 448], ymm14",
            "vmovdqu [{p} + 480], ymm15",
                p = in(reg) r.0.as_mut_ptr(), options(nostack, preserves_flags));
        } else {
            asm!(
            "movdqu [{p} + 0], xmm0",
            "movdqu [{p} + 32], xmm1",
            "movdqu [{p} + 64], xmm2",
            "movdqu [{p} + 96], xmm3",
            "movdqu [{p} + 128], xmm4",
            "movdqu [{p} + 160], xmm5",
            "movdqu [{p} + 192], xmm6",
            "movdqu [{p} + 224], xmm7",
            "movdqu [{p} + 256], xmm8",
            "movdqu [{p} + 288], xmm9",
            "movdqu [{p} + 320], xmm10",
            "movdqu [{p} + 352], xmm11",
            "movdqu [{p} + 384], xmm12",
            "movdqu [{p} + 416], xmm13",
            "movdqu [{p} + 448], xmm14",
            "movdqu [{p} + 480], xmm15",
                p = in(reg) r.0.as_mut_ptr(), options(nostack, preserves_flags));
        }
    }
    r
}

fn get_mxcsr() -> u32 {
    let mut m = 0u32;
    unsafe { asm!("stmxcsr [{}]", in(reg) &mut m, options(nostack, preserves_flags)) };
    m
}
fn set_mxcsr(m: u32) {
    unsafe { asm!("ldmxcsr [{}]", in(reg) &m, options(nostack, readonly, preserves_flags)) };
}

/// A register image no other `seed` produces (only the parts this CPU has).
fn pattern(seed: u64) -> Regs {
    let mut r = Regs([0; 512]);
    let mut x = seed.wrapping_mul(0x9e37_79b9_7f4a_7c15) | 1;
    for (i, chunk) in r.0.chunks_mut(8).enumerate() {
        x ^= x << 13;
        x ^= x >> 7;
        x ^= x << 17;
        if avx() || i % 4 < 2 {
            chunk.copy_from_slice(&x.to_le_bytes());
        }
    }
    r
}
/// A distinct MXCSR per seed: all exceptions masked, rounding mode (bits
/// 13-14) and flush-to-zero / denormals-are-zero (15, 6) from the seed.
fn mxcsr_for(seed: u64) -> u32 {
    0x1f80 | ((seed as u32 & 3) << 13) | ((seed as u32 >> 2 & 1) << 15) | ((seed as u32 >> 3 & 1) << 6)
}

/// Load a pattern, spin and yield across many timer ticks (so the thread
/// is switched out and probably migrated), checking it now and then.
/// Returns how many checks found a wrong register.
fn hammer(seed: u64, rounds: u32) -> u32 {
    let want = pattern(seed);
    let mx = mxcsr_for(seed);
    load(&want);
    set_mxcsr(mx);
    let mut bad = 0;
    for round in 0..rounds {
        let mut x = round as u64;
        for _ in 0..20_000 {
            x = core::hint::black_box(x.wrapping_mul(6364136223846793005).wrapping_add(1));
        }
        if round % 3 == 0 {
            sched_yield();
        }
        if store() != want || get_mxcsr() != mx {
            bad += 1;
            load(&want); // keep going with good values
            set_mxcsr(mx);
        }
    }
    set_mxcsr(0x1f80);
    bad
}

// --- signal handlers -----------------------------------------------------------

/// Offsets in the ucontext the kernel passes (same as Linux's): the
/// `fpstate` pointer in uc_mcontext, and inside the FXSAVE area: MXCSR,
/// xmm0, the software-reserved bytes (sw_bytes) and, in the XSAVE area,
/// the upper halves of ymm0-15.
const UC_FLAGS: usize = 0;
const UC_FPSTATE: usize = 40 + 23 * 8;
const FX_MXCSR: usize = 24;
const FX_XMM: usize = 160;
const FX_SW: usize = 464;
const XS_YMM_HI: usize = 576;
const MAGIC1: u32 = 0x4650_5853;
const MAGIC2: u32 = 0x4650_5845;

static SYNC_SEED: AtomicU64 = AtomicU64::new(0);
/// Bit flags the handler sets for main to check.
static H_CLEAN: AtomicU32 = AtomicU32::new(0); // registers zero + MXCSR default on entry
static H_SAVED: AtomicU32 = AtomicU32::new(0); // the frame held the interrupted values
static H_LAYOUT: AtomicU32 = AtomicU32::new(0); // magic numbers / flags right
static ASYNC_HITS: AtomicU32 = AtomicU32::new(0);
const NEW_XMM15: u64 = 0x1515_1515_cafe_f00d;

fn rd32(p: *const u8) -> u32 {
    unsafe { core::ptr::read_unaligned(p as *const u32) }
}
fn rd64(p: *const u8) -> u64 {
    unsafe { core::ptr::read_unaligned(p as *const u64) }
}

/// Scribble over every vector register and MXCSR (a handler that does
/// float maths, or calls memcpy, would).
fn clobber() {
    load(&pattern(0xdead));
    set_mxcsr(0x1f80 | (3 << 13));
}

extern "C" fn on_usr1(_sig: i32, _info: *const SigInfo, uc: *mut u8) {
    let clean = store() == Regs([0; 512]) && get_mxcsr() == 0x1f80;
    H_CLEAN.store(clean as u32, Ordering::SeqCst);
    let want = pattern(SYNC_SEED.load(Ordering::SeqCst));
    let fp = rd64(unsafe { uc.add(UC_FPSTATE) }) as *mut u8;
    if fp.is_null() || fp as u64 % 64 != 0 {
        H_SAVED.store(0, Ordering::SeqCst);
    } else {
        // xmm i is at FX_XMM + 16*i; our image has register i at 32*i.
        let mut ok = rd32(unsafe { fp.add(FX_MXCSR) }) == mxcsr_for(SYNC_SEED.load(Ordering::SeqCst));
        for i in 0..16 {
            for b in 0..16 {
                ok &= unsafe { *fp.add(FX_XMM + 16 * i + b) } == want.0[32 * i + b];
            }
        }
        let flags = rd64(unsafe { uc.add(UC_FLAGS) });
        let magic1 = rd32(unsafe { fp.add(FX_SW) });
        let layout = if magic1 == MAGIC1 {
            // XSAVE frame: sw_bytes say how big it is; magic2 at the end.
            let size = rd32(unsafe { fp.add(FX_SW + 16) }) as usize;
            let xfeatures = rd64(unsafe { fp.add(FX_SW + 8) });
            if avx() {
                for i in 0..16 {
                    for b in 0..16 {
                        ok &= unsafe { *fp.add(XS_YMM_HI + 16 * i + b) } == want.0[32 * i + 16 + b];
                    }
                }
            }
            flags & 1 != 0 && rd32(unsafe { fp.add(size) }) == MAGIC2 && xfeatures & 3 == 3 && (!avx() || xfeatures & 4 != 0)
        } else {
            // FXSAVE-only frame (no XSAVE): fine as long as we're not using AVX.
            magic1 == 0 && !avx()
        };
        H_SAVED.store(ok as u32, Ordering::SeqCst);
        H_LAYOUT.store(layout as u32, Ordering::SeqCst);
        // Change the saved xmm15: sigreturn must load what's in the frame.
        unsafe { core::ptr::write_unaligned(fp.add(FX_XMM + 16 * 15) as *mut u64, NEW_XMM15) };
    }
    clobber();
}

extern "C" fn on_usr2(_sig: i32) {
    clobber();
    ASYNC_HITS.fetch_add(1, Ordering::SeqCst);
}

static FPE_CODE: AtomicU32 = AtomicU32::new(0);
/// SIGFPE from an unmasked SSE exception. Returning would just run the
/// faulting instruction again, so mask all exceptions (and clear the
/// flags) in the *saved* MXCSR first: the retry then gives inf quietly.
extern "C" fn on_fpe(_sig: i32, info: *const SigInfo, uc: *mut u8) {
    FPE_CODE.store(unsafe { (*info).code } as u32, Ordering::SeqCst);
    let fp = rd64(unsafe { uc.add(UC_FPSTATE) }) as *mut u8;
    if !fp.is_null() {
        unsafe { core::ptr::write_unaligned(fp.add(FX_MXCSR) as *mut u32, 0x1f80) };
    }
}

// --- the tests ---------------------------------------------------------------------

/// N threads in this process, each with its own seed; returns bad checks.
fn thread_round(base: u64, n: u64, rounds: u32) -> u32 {
    let bad = Arc::new(AtomicU32::new(0));
    let mut ts = Vec::new();
    for i in 0..n {
        let bad = bad.clone();
        ts.push(spawn_thread(0, move || {
            bad.fetch_add(hammer(base + i, rounds), Ordering::SeqCst);
        }));
    }
    let main_bad = hammer(base + 99, rounds);
    for t in ts {
        match t {
            Ok(t) => t.join(),
            Err(_) => return 1000,
        }
    }
    bad.load(Ordering::SeqCst) + main_bad
}

fn exec_check() -> i64 {
    // Run straight after exec: nothing has touched the vector registers.
    let clean = store() == Regs([0; 512]) && get_mxcsr() == 0x1f80;
    if clean { 0 } else { 1 }
}

fn main() -> i64 {
    AVX.store(detect_avx(), Ordering::Relaxed);
    if arg(1) == Some("exec-check") {
        return exec_check();
    }
    let mut c = Checks::new("fputest");
    println!(
        "fputest: {} registers, {} CPU(s)",
        if avx() { "ymm0-15 (AVX)" } else { "xmm0-15 (SSE)" },
        sysinfo().cpus
    );

    // 1. The registers start clean (main ran just after exec).
    c.check(exec_check() == 0, "a new program starts with zeroed registers and MXCSR 0x1f80");

    // 2. Threads and processes, all with distinct values, on all CPUs.
    let mut children = Vec::new();
    for p in 1..=2u64 {
        let pid = fork();
        if pid == 0 {
            exit(thread_round(p * 1000, 3, 60).min(100) as i64);
        }
        children.push(pid);
    }
    let bad = thread_round(0, 3, 60);
    let mut child_bad = 0;
    for pid in children {
        child_bad += wait(pid);
    }
    c.check(bad == 0, "4 threads kept their own vector registers + MXCSR");
    c.check(child_bad == 0, "so did 4 threads in each of 2 forked processes (12 at once)");

    // 3. fork copies the registers; exec clears them.
    let want = pattern(7);
    load(&want);
    set_mxcsr(mxcsr_for(7));
    let pid = fork();
    if pid == 0 {
        let same = store() == want && get_mxcsr() == mxcsr_for(7);
        if !same {
            exit(2);
        }
        execve("/mnt/disk0/bin/fputest", &["fputest", "exec-check"], &[]);
        exit(3);
    }
    let parent_same = store() == want;
    set_mxcsr(0x1f80);
    let code = wait(pid);
    c.check(parent_same && code != 2, "fork: the child starts with the parent's registers");
    c.check(code == 0, "exec: the new program starts with clean registers");

    // 4. A synchronous signal whose handler clobbers everything.
    sigaction_info(SIGUSR1, on_usr1, 0);
    SYNC_SEED.store(42, Ordering::SeqCst);
    let mut want = pattern(42);
    load(&want);
    set_mxcsr(mxcsr_for(42));
    kill(getpid(), SIGUSR1); // the handler runs on the way back from this
    let got = store();
    let mx = get_mxcsr();
    set_mxcsr(0x1f80);
    c.check(H_CLEAN.load(Ordering::SeqCst) == 1, "the handler starts with clean registers");
    c.check(H_SAVED.load(Ordering::SeqCst) == 1, "uc_mcontext.fpstate holds the interrupted registers");
    c.check(H_LAYOUT.load(Ordering::SeqCst) == 1, "the fpstate has Linux's layout (sw_bytes, magics, uc_flags)");
    want.0[32 * 15..32 * 15 + 8].copy_from_slice(&NEW_XMM15.to_le_bytes());
    c.check(got == want && mx == mxcsr_for(42), "after sigreturn: our registers back, with the handler's xmm15 edit");

    // 5. Asynchronous signals: another thread interrupts us while we spin.
    signal(SIGUSR2, on_usr2);
    let me = gettid();
    let sender = spawn_thread(0, move || {
        for _ in 0..40 {
            tgkill(getpid(), me, SIGUSR2);
            sleep_ms(2);
        }
    });
    let want = pattern(4242);
    load(&want);
    set_mxcsr(mxcsr_for(4242));
    let t0 = now_us();
    let mut bad = 0;
    while ASYNC_HITS.load(Ordering::SeqCst) < 40 && now_us() - t0 < 5_000_000 {
        if store() != want || get_mxcsr() != mxcsr_for(4242) {
            bad += 1;
            load(&want);
            set_mxcsr(mxcsr_for(4242));
        }
    }
    set_mxcsr(0x1f80);
    if let Ok(t) = sender {
        t.join();
    }
    let hits = ASYNC_HITS.load(Ordering::SeqCst);
    println!("fputest: {} async signals, each handler clobbered every register", hits);
    c.check(hits == 40 && bad == 0, "40 async signals with clobbering handlers changed nothing we saw");

    // 6. An unmasked SSE exception (divide by zero) -> SIGFPE, FPE_FLTDIV.
    sigaction_info(SIGFPE, on_fpe, 0);
    let one = 1.0f64.to_bits();
    let zero = 0u64;
    let mut out = 0u64;
    set_mxcsr(0x1f80 & !(1 << 9)); // unmask ZE (bit 9)
    unsafe {
        asm!(
            "movq xmm0, [{a}]",
            "divsd xmm0, [{b}]", // 1.0 / 0.0
            "movq [{o}], xmm0",
            a = in(reg) &one, b = in(reg) &zero, o = in(reg) &mut out,
            options(nostack)
        );
    }
    let mx = get_mxcsr();
    set_mxcsr(0x1f80);
    let code = FPE_CODE.load(Ordering::SeqCst);
    if code == 0 && out == f64::INFINITY.to_bits() && mx & 4 != 0 {
        // QEMU's software CPU (TCG) sets the flag but never raises #XM.
        println!("fputest: skipped SIGFPE check: this CPU doesn't raise SIMD exceptions (QEMU TCG)");
    } else {
        c.check(code == 3, "unmasked 1.0/0.0 -> SIGFPE with si_code FPE_FLTDIV");
        c.check(out == f64::INFINITY.to_bits() && mx & (1 << 9) != 0, "the handler's MXCSR edit took effect (retry gave inf)");
    }

    c.done()
}