smarm/src/context.rs

//! Cooperative context switching, x86-64.
//!
//! Two naked-asm functions move execution between a scheduler thread and an
//! actor running on its own mmap'd stack. The compiler cannot do this; the
//! whole point of `#[unsafe(naked)]` is that we control every instruction.
//!
//! `SCHEDULER_SP` and `ACTOR_SP` are thread-locals holding each side's saved
//! stack pointer. `init_actor_stack` builds the initial stack so that the
//! first `switch_to_actor` lands inside the entry function with `rsp % 16 == 8`
//! (the x86-64 ABI requirement at function entry).

use std::cell::Cell;

thread_local! {
    static SCHEDULER_SP: Cell<usize> = const { Cell::new(0) };
    static ACTOR_SP:     Cell<usize> = const { Cell::new(0) };
}

fn  get_scheduler_sp() -> usize     { SCHEDULER_SP.with(|c| c.get()) }
fn  set_scheduler_sp(v: usize)      { SCHEDULER_SP.with(|c| c.set(v)) }
pub fn get_actor_sp() -> usize      { ACTOR_SP.with(|c| c.get()) }
pub fn set_actor_sp(v: usize)       { ACTOR_SP.with(|c| c.set(v)) }

// ---------------------------------------------------------------------------
// Initial stack layout
//
// We start from aligned_top = top & ~15, then bias by 8 and push seven 8-byte
// slots (downward). The first `switch_to_actor` pops r15..rbx and `ret`s —
// landing in `entry` with rsp % 16 == 8 (the x86-64 ABI state at the point
// just after a `call`, which is what `entry` is compiled to expect).
//
// Layout (high → low), relative to aligned_top = top & ~15. Note the entry
// slot is at aligned_top - 16, NOT aligned_top - 8: the function does
// `(top & ~15) - 8` and *then* a `-= 8` before the first write, so the first
// stored word lands at aligned_top - 16. Verified by single-stepping the
// `ret` under llmdbg: entry sits at an address with %16 == 0, so the post-ret
// rsp is %16 == 8.
//
//   aligned_top - 8  : (unused padding; keeps entry's slot %16 == 0)
//   aligned_top - 16 : entry ptr       ← `ret` target. Post-ret: rsp % 16 == 8.
//   aligned_top - 24 : rbx = 0
//   aligned_top - 32 : rbp = 0
//   aligned_top - 40 : r12 = 0
//   aligned_top - 48 : r13 = 0
//   aligned_top - 56 : r14 = 0
//   aligned_top - 64 : r15 = 0         ← initial rsp
// ---------------------------------------------------------------------------

pub fn init_actor_stack(top: *mut u8, entry: extern "C-unwind" fn()) -> usize {
    unsafe {
        let mut sp = (top as usize & !15) - 8;
        sp -= 8; (sp as *mut usize).write(entry as usize); // ret target
        sp -= 8; (sp as *mut usize).write(0); // rbx
        sp -= 8; (sp as *mut usize).write(0); // rbp
        sp -= 8; (sp as *mut usize).write(0); // r12
        sp -= 8; (sp as *mut usize).write(0); // r13
        sp -= 8; (sp as *mut usize).write(0); // r14
        sp -= 8; (sp as *mut usize).write(0); // r15
        sp
    }
}

// ---------------------------------------------------------------------------
// Context switch shims
//
// Each shim:
//   1. Pushes the six callee-saved integer registers.
//   2. Snaps rsp into rdi and calls the Rust helper that stores it.
//   3. Calls the Rust helper that returns the *other* side's saved rsp.
//   4. Moves that into rsp.
//   5. Pops the six registers and rets.
//
// XMM registers are NOT saved here. We rely on every yield happening through
// a Rust call site, which means the compiler has spilled any live XMM state
// to the stack before we get here. (This is the same argument the compiler
// uses internally — callee-saved regs are what survive a `call`, and the
// SysV AMD64 ABI says XMM0–15 are all caller-saved.) If we ever yield from
// a place that isn't a Rust call boundary, this assumption breaks.
// ---------------------------------------------------------------------------

#[unsafe(naked)]
unsafe extern "C" fn switch_to_actor_asm() {
    core::arch::naked_asm!(
        "push rbx", "push rbp", "push r12", "push r13", "push r14", "push r15",
        "mov rdi, rsp",
        "call {set_sched_sp}",
        "call {get_actor_sp}",
        "mov rsp, rax",
        "pop r15", "pop r14", "pop r13", "pop r12", "pop rbp", "pop rbx",
        "ret",
        set_sched_sp = sym set_scheduler_sp,
        get_actor_sp = sym get_actor_sp,
    );
}

/// Resume the actor whose sp is in `ACTOR_SP`. Returns when the actor yields.
///
/// # Safety
///
/// The caller must be running on a scheduler thread with a valid actor stack
/// pointer installed in `ACTOR_SP` — either by `init_actor_stack` (first
/// resume) or by a prior `switch_to_scheduler` (subsequent resumes). Resuming
/// with an unset or stale `ACTOR_SP` transfers control to an arbitrary address.
/// Must not be called from within an actor (only the scheduler side may resume).
pub unsafe fn switch_to_actor() {
    unsafe { switch_to_actor_asm() };
}

/// Yield from the running actor back to its scheduler thread. Returns when the
/// actor is next resumed via [`switch_to_actor`].
///
/// # Safety
///
/// The caller must be running on an actor stack that was entered through
/// [`switch_to_actor`], so that `SCHEDULER_SP` holds the live saved stack
/// pointer of the scheduler side. Calling this from the scheduler thread, or
/// before any actor has been resumed, transfers control to an arbitrary
/// address.
#[unsafe(naked)]
pub unsafe extern "C" fn switch_to_scheduler() {
    core::arch::naked_asm!(
        "push rbx", "push rbp", "push r12", "push r13", "push r14", "push r15",
        "mov rdi, rsp",
        "call {set_actor_sp}",
        "call {get_sched_sp}",
        "mov rsp, rax",
        "pop r15", "pop r14", "pop r13", "pop r12", "pop rbp", "pop rbx",
        "ret",
        set_actor_sp = sym set_actor_sp,
        get_sched_sp = sym get_scheduler_sp,
    );
}