smarm-dev
6d17254ae3
The diagram claimed the entry/ret target sits at aligned_top-8 and r15 at aligned_top-56. The code does (top & ~15) - 8 then a -= 8 before the first write, so entry actually lands at aligned_top-16 and r15 at aligned_top-64. Confirmed by single-stepping the shim's ret under llmdbg: the entry slot is at an address with %16==0, giving the required rsp%16==8 on entry. Code was correct; only the comment was wrong (and would mislead a maintainer).
115 lines
4.7 KiB
Rust
115 lines
4.7 KiB
Rust
//! 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.
|
||
pub unsafe fn switch_to_actor() {
|
||
unsafe { switch_to_actor_asm() };
|
||
}
|
||
|
||
#[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,
|
||
);
|
||
}
|