//! Reschedule and leave-root IPIs, serviced without the kernel lock.

use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

use super::sbi::sbi_send_ipi;
use super::{MAX_HARTS, hart_id, kernel_root, num_harts};
use crate::mmu;

const SIP_SSIP: usize = 1 << 1;

/// The root table each hart translates with.
static ACTIVE_ROOT: [AtomicUsize; MAX_HARTS] = [const { AtomicUsize::new(0) }; MAX_HARTS];
static RESCHEDULE: [AtomicBool; MAX_HARTS] = [const { AtomicBool::new(false) }; MAX_HARTS];
static LEAVE_ROOT: [AtomicUsize; MAX_HARTS] = [const { AtomicUsize::new(0) }; MAX_HARTS];
static LEFT: [AtomicUsize; MAX_HARTS] = [const { AtomicUsize::new(0) }; MAX_HARTS];

/// Called by `mmu::activate`.
pub fn note_active_root(root: usize) {
    if let Some(r) = ACTIVE_ROOT.get(hart_id()) {
        r.store(root, Ordering::SeqCst);
    }
}

/// Whether this hart has an IPI or a reschedule request it has not acted
/// on; fast paths must then take the slow path.
pub fn has_pending_request() -> bool {
    let sip: usize;
    // SAFETY: reading sip has no side effects.
    unsafe { core::arch::asm!("csrr {0}, sip", out(reg) sip) };
    sip & SIP_SSIP != 0
        || RESCHEDULE
            .get(hart_id())
            .is_some_and(|r| r.load(Ordering::SeqCst))
}

// ---- IPIs ----

/// The lock-free half of an IPI: acknowledges it and leaves the root being
/// changed, if asked. Safe to call at any time.
pub fn service_ipi() {
    let sip: usize;
    // SAFETY: reading sip has no side effects.
    unsafe { core::arch::asm!("csrr {0}, sip", out(reg) sip) };
    if sip & SIP_SSIP == 0 {
        return;
    }
    // Clear first: a request posted after this raises SSIP again.
    // SAFETY: clearing our own pending software interrupt.
    unsafe { core::arch::asm!("csrc sip, {0}", in(reg) SIP_SSIP) };
    let me = hart_id();
    let (Some(leave), Some(left)) = (LEAVE_ROOT.get(me), LEFT.get(me)) else {
        return;
    };
    let root = leave.swap(0, Ordering::SeqCst);
    if root != 0 {
        if ACTIVE_ROOT.get(me).map(|r| r.load(Ordering::SeqCst)) == Some(root) {
            mmu::activate(mmu::page_table::table_at(kernel_root()));
        }
        // SAFETY: flushing translations has no other effect.
        unsafe { core::arch::asm!("sfence.vma zero, zero") };
        left.fetch_add(1, Ordering::SeqCst);
    }
}

/// Asks `hart` to look at its run queue again, on its way out of the
/// kernel.
pub fn kick(hart: usize) {
    if let Some(r) = RESCHEDULE.get(hart) {
        r.store(true, Ordering::SeqCst);
        sbi_send_ipi(hart);
    }
}

/// Consumes this hart's reschedule request.
pub fn take_reschedule() -> bool {
    RESCHEDULE
        .get(hart_id())
        .is_some_and(|r| r.swap(false, Ordering::SeqCst))
}

/// Makes every other hart leave `root` and flush, waiting for each.
/// Kernel lock held, so no hart can start using `root` meanwhile.
pub fn shoot_down(root: usize) {
    let me = hart_id();
    for hart in 0..num_harts() {
        if hart == me || ACTIVE_ROOT[hart].load(Ordering::SeqCst) != root {
            continue;
        }
        let seen = LEFT[hart].load(Ordering::SeqCst);
        LEAVE_ROOT[hart].store(root, Ordering::SeqCst);
        sbi_send_ipi(hart);
        while LEFT[hart].load(Ordering::SeqCst) == seen {
            core::hint::spin_loop();
        }
    }
}