vj-sim 0.2.0

//! Stackful fibers for blocking DPI import tasks (LRM 1800-2023 35.5.3,
//! 35.7, 35.9).
//!
//! A DPI `import task` is a C function that may call an exported SV task
//! which consumes time. The C frame cannot be unwound into a stackless
//! process state machine, so the call runs on a [`Fiber`]: a private stack
//! the scheduler can leave ([`yield_now`], from inside the fiber) and later
//! re-enter ([`Fiber::resume`]). A call that never blocks finishes on its
//! first `resume` and costs one stack switch pair.
//!
//! Built on glibc's `makecontext`/`swapcontext` (declared here directly: no
//! `libc` dependency in this workspace). `ucontext_t` is kept as an opaque,
//! over-sized buffer; the only fields touched go through those functions,
//! except `uc_stack` / `uc_link`, which are located by the x86_64 / aarch64
//! glibc layout checked in the tests.
//!
//! A panic never crosses the fiber boundary: the entry closure runs under
//! `catch_unwind` and the payload is re-raised by [`Fiber::resume`] on the
//! resuming stack.

use std::any::Any;
use std::cell::Cell;
use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind};

/// Bytes reserved for a `ucontext_t` (936 on x86_64 glibc, 4560 on aarch64).
const UCONTEXT_BYTES: usize = 5120;
/// Fiber stack size. A whole process kernel can run on a fiber (a design with a
/// timed export task, `CompiledKernelProcess::with_task_fibers`), so it is as
/// deep as the kernels need; the pages are committed only when touched.
const STACK_BYTES: usize = 4 << 20;
/// Stacks kept for reuse after a fiber is dropped (a mapping costs more than a
/// short call), per thread.
const POOL_MAX: usize = 32;
const GUARD_BYTES: usize = 4096;

#[repr(C, align(64))]
struct UContext([u8; UCONTEXT_BYTES]);

unsafe extern "C" {
    fn getcontext(ucp: *mut UContext) -> i32;
    fn makecontext(ucp: *mut UContext, func: extern "C" fn(), argc: i32, ...);
    fn swapcontext(oucp: *mut UContext, ucp: *const UContext) -> i32;
    fn mmap(addr: *mut u8, len: usize, prot: i32, flags: i32, fd: i32, off: i64) -> *mut u8;
    fn munmap(addr: *mut u8, len: usize) -> i32;
    fn mprotect(addr: *mut u8, len: usize, prot: i32) -> i32;
}

const PROT_NONE: i32 = 0;
const PROT_RW: i32 = 3;
const MAP_PRIVATE_ANON_STACK: i32 = 0x02 | 0x20 | 0x20000 | 0x4000; // PRIVATE | ANONYMOUS | STACK | NORESERVE
const MAP_FAILED: *mut u8 = !0usize as *mut u8;

/// Whether a [`Fiber::resume`] left the fiber parked or ran it to its end.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FiberState {
    /// The fiber called [`yield_now`]; resume it later.
    Yielded,
    /// The entry closure returned; the fiber is spent.
    Finished,
}

struct Inner {
    /// The fiber's saved context while it is parked.
    fiber_ctx: UContext,
    /// The resumer's saved context while the fiber runs.
    caller_ctx: UContext,
    entry: Option<Box<dyn FnOnce()>>,
    panic: Option<Box<dyn Any + Send>>,
    done: bool,
    map: *mut u8,
}

/// One suspendable C call.
pub struct Fiber {
    inner: Box<Inner>,
}

thread_local! {
    /// Unused stack mappings (guard page included), most recent last.
    static POOL: std::cell::RefCell<Vec<usize>> = const { std::cell::RefCell::new(Vec::new()) };
    /// The fiber currently running on this thread (null on the main stack).
    static CURRENT: Cell<*mut Inner> = const { Cell::new(std::ptr::null_mut()) };
}

extern "C" fn trampoline() {
    let inner = CURRENT.with(Cell::get);
    // SAFETY: `resume` set CURRENT to the live `Inner` before switching here.
    let f = unsafe { (*inner).entry.take() };
    if let Some(f) = f
        && let Err(p) = catch_unwind(AssertUnwindSafe(f))
    {
        // SAFETY: as above.
        unsafe { (*inner).panic = Some(p) };
    }
    // SAFETY: as above. Returning to the resumer; never resumed again.
    unsafe {
        (*inner).done = true;
        swapcontext(&raw mut (*inner).fiber_ctx, &raw const (*inner).caller_ctx);
    }
    unreachable!("a finished fiber was resumed");
}

impl Fiber {
    /// A fiber that will run `entry` on its own stack on the first
    /// [`Self::resume`]. `None` when the stack cannot be mapped.
    pub fn new(entry: impl FnOnce() + 'static) -> Option<Fiber> {
        let total = STACK_BYTES + GUARD_BYTES;
        let pooled = POOL.with(|p| p.borrow_mut().pop());
        let map = match pooled {
            Some(m) => m as *mut u8,
            None => {
                // SAFETY: plain anonymous mapping; checked for failure.
                let map = unsafe {
                    mmap(
                        std::ptr::null_mut(),
                        total,
                        PROT_RW,
                        MAP_PRIVATE_ANON_STACK,
                        -1,
                        0,
                    )
                };
                if map == MAP_FAILED {
                    return None;
                }
                // SAFETY: the lowest page of the mapping, made a guard.
                unsafe { mprotect(map, GUARD_BYTES, PROT_NONE) };
                map
            }
        };
        let mut inner = Box::new(Inner {
            fiber_ctx: UContext([0; UCONTEXT_BYTES]),
            caller_ctx: UContext([0; UCONTEXT_BYTES]),
            entry: Some(Box::new(entry)),
            panic: None,
            done: false,
            map,
        });
        // SAFETY: `fiber_ctx` is a zeroed, aligned buffer larger than
        // `ucontext_t`; `set_stack` writes the documented stack fields.
        unsafe {
            if getcontext(&raw mut inner.fiber_ctx) != 0 {
                release_stack(map);
                return None;
            }
            set_stack(
                &mut inner.fiber_ctx,
                map.add(GUARD_BYTES),
                STACK_BYTES,
                std::ptr::null_mut(),
            );
            makecontext(&raw mut inner.fiber_ctx, trampoline, 0);
        }
        Some(Fiber { inner })
    }

    /// Runs the fiber until it yields or finishes. Resumes a panic raised
    /// inside it. Must not be called from inside the same fiber.
    pub fn resume(&mut self) -> FiberState {
        assert!(!self.inner.done, "resume of a finished fiber");
        let inner: *mut Inner = &raw mut *self.inner;
        let prev = CURRENT.with(|c| c.replace(inner));
        // SAFETY: both contexts live in the boxed `Inner`.
        unsafe { swapcontext(&raw mut (*inner).caller_ctx, &raw const (*inner).fiber_ctx) };
        CURRENT.with(|c| c.set(prev));
        if let Some(p) = self.inner.panic.take() {
            resume_unwind(p);
        }
        if self.inner.done {
            FiberState::Finished
        } else {
            FiberState::Yielded
        }
    }

    /// Whether the entry closure has returned.
    pub fn is_finished(&self) -> bool {
        self.inner.done
    }
}

impl Drop for Fiber {
    fn drop(&mut self) {
        // A parked fiber's stack is simply discarded: its frames never run
        // their destructors (an abandoned blocked import task, e.g. at the
        // end of simulation). Nothing on it is owned by anyone else.
        release_stack(self.inner.map);
    }
}

/// Returns a stack mapping made in [`Fiber::new`], unused from now on, to the
/// pool (or to the OS when the pool is full).
fn release_stack(map: *mut u8) {
    let kept = POOL.with(|p| {
        let mut p = p.borrow_mut();
        let room = p.len() < POOL_MAX;
        if room {
            p.push(map as usize);
        }
        room
    });
    if !kept {
        // SAFETY: the mapping made in `new`, unused once the fiber is gone.
        unsafe { munmap(map, STACK_BYTES + GUARD_BYTES) };
    }
}

/// Parks the running fiber and returns to its resumer. `false` (nothing
/// happened) when called on the main stack.
pub fn yield_now() -> bool {
    let inner = CURRENT.with(Cell::get);
    if inner.is_null() {
        return false;
    }
    // SAFETY: CURRENT is the live fiber that is calling us.
    unsafe { swapcontext(&raw mut (*inner).fiber_ctx, &raw const (*inner).caller_ctx) };
    true
}

/// Whether the caller runs on a fiber (a blocking export is legal here).
pub fn on_fiber() -> bool {
    !CURRENT.with(Cell::get).is_null()
}

/// Writes `uc_link`, `uc_stack.{ss_sp,ss_size}` of a glibc `ucontext_t`.
///
/// glibc x86_64 and aarch64 share the head layout
/// `{ unsigned long uc_flags; ucontext_t *uc_link; stack_t uc_stack; ... }`
/// with `stack_t = { void *ss_sp; int ss_flags; size_t ss_size; }`.
///
/// # Safety
/// `uc` must have been filled by `getcontext`.
unsafe fn set_stack(uc: &mut UContext, sp: *mut u8, size: usize, link: *mut UContext) {
    let w = uc.0.as_mut_ptr().cast::<usize>();
    // SAFETY: offsets 1..=4 words lie within the buffer.
    unsafe {
        w.add(1).write(link as usize); // uc_link
        w.add(2).write(sp as usize); // ss_sp
        w.add(3).write(0); // ss_flags (+ padding)
        w.add(4).write(size); // ss_size
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::rc::Rc;

    #[test]
    fn runs_to_completion_without_yielding() {
        let hit = Rc::new(Cell::new(0));
        let h = hit.clone();
        let mut f = Fiber::new(move || h.set(7)).unwrap();
        assert!(!on_fiber());
        assert_eq!(f.resume(), FiberState::Finished);
        assert_eq!(hit.get(), 7);
        assert!(f.is_finished());
    }

    #[test]
    fn yields_and_resumes_in_order() {
        let log = Rc::new(std::cell::RefCell::new(Vec::new()));
        let l = log.clone();
        let mut f = Fiber::new(move || {
            assert!(on_fiber());
            l.borrow_mut().push(1);
            assert!(yield_now());
            l.borrow_mut().push(3);
            yield_now();
            l.borrow_mut().push(5);
        })
        .unwrap();
        assert_eq!(f.resume(), FiberState::Yielded);
        log.borrow_mut().push(2);
        assert_eq!(f.resume(), FiberState::Yielded);
        log.borrow_mut().push(4);
        assert_eq!(f.resume(), FiberState::Finished);
        assert_eq!(*log.borrow(), [1, 2, 3, 4, 5]);
    }

    #[test]
    fn interleaved_fibers_keep_their_own_stacks() {
        let out = Rc::new(std::cell::RefCell::new(Vec::new()));
        let mk = |tag: i32| {
            let o = out.clone();
            Fiber::new(move || {
                let local = [tag; 64];
                yield_now();
                o.borrow_mut().push(local.iter().sum::<i32>());
            })
            .unwrap()
        };
        let (mut a, mut b) = (mk(1), mk(2));
        a.resume();
        b.resume();
        // Out-of-order completion: the case a nested scheduler cannot do.
        assert_eq!(b.resume(), FiberState::Finished);
        assert_eq!(a.resume(), FiberState::Finished);
        assert_eq!(*out.borrow(), [128, 64]);
    }

    #[test]
    fn yield_on_main_stack_is_a_no_op() {
        assert!(!yield_now());
    }

    #[test]
    fn panic_is_re_raised_on_the_resumer() {
        let mut f = Fiber::new(|| panic!("boom")).unwrap();
        let r = catch_unwind(AssertUnwindSafe(|| f.resume()));
        assert!(r.is_err());
    }

    #[test]
    fn nested_fibers() {
        let log = Rc::new(std::cell::RefCell::new(Vec::new()));
        let l = log.clone();
        let mut outer = Fiber::new(move || {
            let l2 = l.clone();
            let mut inner = Fiber::new(move || {
                l2.borrow_mut().push("in1");
                yield_now();
                l2.borrow_mut().push("in2");
            })
            .unwrap();
            inner.resume();
            l.borrow_mut().push("out1");
            yield_now();
            inner.resume();
            l.borrow_mut().push("out2");
        })
        .unwrap();
        outer.resume();
        outer.resume();
        assert_eq!(*log.borrow(), ["in1", "out1", "in2", "out2"]);
    }
}