//! SMP self-test, run when the disk carries `bin/smptest`: after revoking a
//! spinner's page or killing it, no spinner counter may move on any hart.

use core::sync::atomic::{AtomicUsize, Ordering};
use libkos::fs::Storage;
use libkos::syscall;
use types::address::VirtualAddress;
use types::bounded::{CapSlot, PageCount};
use types::mapping::PagePerms;
use types::permission::{Permission, PermissionSet};

const OBSERVE_VA: VirtualAddress = VirtualAddress::from_raw(0x0140_0000);
const THREADS: usize = 4;
/// QEMU virt's timebase is 10 MHz.
const TICKS_PER_MS: u64 = 10_000;

/// The spinner's page as init sees it: word 0 hands out counters, then
/// one counter per spinner thread.
type Words = &'static [AtomicUsize];

fn counters(words: Words) -> [usize; THREADS] {
    core::array::from_fn(|i| words[1 + i].load(Ordering::SeqCst))
}

/// Spins (yielding) until `done` or `ms` milliseconds pass.
fn wait_until(ms: u64, mut done: impl FnMut() -> bool) -> bool {
    let deadline = syscall::get_time().raw() + ms * TICKS_PER_MS;
    while syscall::get_time().raw() < deadline {
        if done() {
            return true;
        }
        syscall::yield_now();
    }
    done()
}

/// Starts a spinner on `page`; its process capability.
fn start_spinner(storage: &mut Storage, scratch: &mut [u8], page: CapSlot, words: Words) -> Result<CapSlot, &'static str> {
    for w in words {
        w.store(0, Ordering::SeqCst);
    }
    let spinner = crate::load::load_binary(storage, scratch, b"bin/smptest", crate::manifest::SMALL).ok_or("cannot load bin/smptest")?;
    let rw = PermissionSet::empty()
        .with(Permission::Read)
        .with(Permission::Write)
        .with(Permission::Delegate);
    // The page first: the spinner expects it in slot 0.
    if syscall::delegate(page, spinner.process_cap, rw).is_none() {
        return Err("cannot give the spinner its page");
    }
    crate::launch::start(&spinner, None);
    if !wait_until(3000, || counters(words).iter().all(|c| *c > 0)) {
        return Err("spinner threads did not all start");
    }
    Ok(spinner.process_cap)
}

/// After `stop`, no counter may move.
fn check_stopped(words: Words, what: &'static str) -> Result<(), &'static str> {
    let before = counters(words);
    wait_until(200, || false);
    if counters(words) != before {
        return Err(what);
    }
    Ok(())
}

fn run_phases(storage: &mut Storage, scratch: &mut [u8]) -> Result<(), &'static str> {
    let page = syscall::create_region(PageCount::new(1).ok_or("page count")?).map_err(|_| "no counter page")?;
    let window = libkos::window::map(page, OBSERVE_VA, PagePerms::ReadWrite).map_err(|_| "cannot map the counter page")?;
    let words = window.atomics(types::address::ByteOffset::ZERO, 1 + THREADS);

    let events = syscall::create_notification().map_err(|_| "no notification")?;
    timer_is_on_time(events)?;

    // Revoke must unmap the page on every hart; the spinner then faults and dies.
    let spinner = start_spinner(storage, scratch, page, words)?;
    syscall::watch_death(spinner, events).map_err(|_| "cannot watch the spinner")?;
    if !syscall::revoke(page) {
        return Err("revoke failed");
    }
    check_stopped(words, "a counter moved after its page was revoked (stale TLB)")?;
    if syscall::kill(spinner) {
        return Err("the spinner survived losing its page");
    }
    if syscall::poll(events) != Ok(1) {
        return Err("nobody heard the spinner fault");
    }
    libkos::println!("smptest: revoke stops every hart: ok");

    // Kill must stop threads running on other harts before it returns.
    let spinner = start_spinner(storage, scratch, page, words)?;
    syscall::watch_death(spinner, events).map_err(|_| "cannot watch the spinner")?;
    if !syscall::kill(spinner) {
        return Err("kill failed");
    }
    check_stopped(words, "a counter moved after its process was killed")?;
    if syscall::poll(events) != Ok(1) {
        return Err("nobody heard the spinner killed");
    }
    libkos::println!("smptest: kill stops every hart: ok");
    Ok(())
}

/// A timer must wake its waiter at the deadline, not at the next scheduling
/// tick (a second away).
fn timer_is_on_time(events: CapSlot) -> Result<(), &'static str> {
    let start = syscall::get_time().raw();
    let deadline = start + 50 * TICKS_PER_MS;
    syscall::set_timer(events, types::time::Ticks::from_raw(deadline)).map_err(|_| "set_timer failed")?;
    syscall::wait(events).map_err(|_| "wait failed")?;
    let woke = syscall::get_time().raw();
    if woke < deadline {
        return Err("the timer fired early");
    }
    if woke - deadline > 20 * TICKS_PER_MS {
        return Err("the timer fired late");
    }
    libkos::println!("smptest: timer wakes on time: ok");
    Ok(())
}

/// Runs the hart checks; on failure reports and halts the machine, else
/// returns so the spawn test can run (and report the whole result).
pub fn run(storage: &mut Storage, scratch: &mut [u8]) {
    let Err(why) = run_phases(storage, scratch) else {
        return;
    };
    libkos::println!("SMPTEST FAIL: {}", why);
    match crate::devices::system_control() {
        Some(control) => syscall::halt(control),
        None => {
            libkos::println!("smptest: no SystemControl capability to halt with");
            crate::idle()
        }
    }
}