//! Kernel stacks with guard pages.
//!
//! Every kernel thread (and every CPU's double-fault handler) needs its own
//! stack. Stack overflows are nasty: the stack grows *down* into whatever
//! lies below it and silently corrupts it. So each stack gets a **guard
//! page** underneath: a page that is deliberately never mapped. Overflowing
//! the stack touches the guard page, the CPU raises a page fault (and,
//! because pushing the fault frame onto the broken stack fails too, a double
//! fault on a separate emergency stack), and interrupts.rs prints a clear
//! "kernel stack overflow" message and kills the thread.
//!
//! All stacks live in one 512 GiB region (L4 slot 200), cut into slots:
//!
//! ```text
//!   slot n:  [ guard page (unmapped) ][ 16 pages = 64 KiB of stack ]
//!            ^ STACKS_START + n * SLOT_SIZE                         ^ top
//! ```
//!
//! **Reusing a slot safely on several CPUs.** When a stack is freed we unmap
//! its pages and give the frames back. Another CPU may still have the old
//! translations cached in its TLB (translation lookaside buffer). If the
//! slot got a new stack and a thread ran on that other CPU, it could use the
//! stale entries and write into frames that now belong to someone else.
//! Real kernels send a "TLB shootdown" interrupt to every CPU. We use a
//! simpler trick: the scheduler remembers, per thread, which CPUs have
//! flushed that thread's stack from their TLB, and flushes it (`invlpg`,
//! 16 instructions) the first time the thread runs on a CPU. Nothing else
//! ever touches another thread's stack, so that's enough.

use crate::memory;
use alloc::vec::Vec;
use core::sync::atomic::{AtomicU64, Ordering};
use spin::Mutex;
use x86_64::instructions::interrupts::without_interrupts;
use x86_64::{VirtAddr, structures::paging::Page};

/// Usable size of every kernel stack.
pub const STACK_SIZE: u64 = 64 * 1024;
const PAGE: u64 = 4096;
/// One guard page plus the stack.
const SLOT_SIZE: u64 = PAGE + STACK_SIZE;
/// Start of the stack region (L4 slot 200 = 200 * 512 GiB).
pub const STACKS_START: u64 = (memory::STACKS_L4_INDEX as u64) << 39;
const STACKS_END: u64 = STACKS_START + (1 << 39);

/// Next never-used slot number.
static NEXT_SLOT: AtomicU64 = AtomicU64::new(0);
/// Slots whose stacks were freed, ready for reuse (taken with interrupts off).
static FREE_SLOTS: Mutex<Vec<u64>> = Mutex::new(Vec::new());
/// How many stacks exist right now (for statistics and tests).
static LIVE: AtomicU64 = AtomicU64::new(0);

/// An owned, mapped kernel stack. Dropping it unmaps the pages and frees
/// the frames.
pub struct KernelStack {
    slot: u64,
}

impl KernelStack {
    /// A new 64 KiB stack (zeroed), or None if we're out of memory.
    pub fn new() -> Option<KernelStack> {
        let slot = without_interrupts(|| FREE_SLOTS.lock().pop())
            .unwrap_or_else(|| NEXT_SLOT.fetch_add(1, Ordering::Relaxed));
        let stack = KernelStack { slot };
        LIVE.fetch_add(1, Ordering::Relaxed);
        assert!(stack.top() <= STACKS_END, "kernel stack area exhausted");
        for page in stack.pages() {
            if memory::map_kernel_page(page).is_none() {
                return None; // `stack` is dropped here: what we mapped is undone
            }
        }
        Some(stack)
    }

    /// Lowest usable address (just above the guard page).
    pub fn bottom(&self) -> u64 {
        STACKS_START + self.slot * SLOT_SIZE + PAGE
    }

    /// One past the highest usable address: the initial stack pointer.
    pub fn top(&self) -> u64 {
        self.bottom() + STACK_SIZE
    }

    /// Address of the guard page under this stack.
    pub fn guard_page(&self) -> u64 {
        self.bottom() - PAGE
    }

    fn pages(&self) -> impl Iterator<Item = Page> {
        let bottom = self.bottom();
        (0..STACK_SIZE / PAGE).map(move |i| Page::containing_address(VirtAddr::new(bottom + i * PAGE)))
    }

    /// Forget any TLB entries *this* CPU might hold for this stack's
    /// addresses (left over from an older stack in the same slot).
    pub fn flush_local_tlb(&self) {
        for page in self.pages() {
            x86_64::instructions::tlb::flush(page.start_address());
        }
    }
}

impl Drop for KernelStack {
    fn drop(&mut self) {
        for page in self.pages() {
            memory::unmap_kernel_page(page); // skips pages that never got mapped
        }
        LIVE.fetch_sub(1, Ordering::Relaxed);
        let slot = self.slot;
        without_interrupts(|| FREE_SLOTS.lock().push(slot));
    }
}

/// Is `addr` inside the guard page of some kernel stack?
pub fn is_guard_page(addr: u64) -> bool {
    (STACKS_START..STACKS_END).contains(&addr) && (addr - STACKS_START) % SLOT_SIZE < PAGE
}

/// How many kernel stacks are allocated right now.
pub fn live_stacks() -> u64 {
    LIVE.load(Ordering::Relaxed)
}