smarm/tests/context.rs

//! Low-level context-switch tests. These poke `init_actor_stack` and the
//! naked asm shims directly — no scheduler involved.

use smarm::context::{
    get_actor_sp, init_actor_stack, set_actor_sp, switch_to_actor, switch_to_scheduler,
};
use smarm::stack::Stack;
use std::cell::Cell;

thread_local! {
    static LOG: Cell<u64> = const { Cell::new(0) };
}

fn log(v: u64) { LOG.with(|c| c.set(c.get() | v)); }
fn get_log() -> u64 { LOG.with(|c| c.get()) }
fn reset_log() { LOG.with(|c| c.set(0)); }

extern "C-unwind" fn actor_simple() {
    log(0x1);
    unsafe { switch_to_scheduler() };
}

#[test]
fn actor_runs_and_returns_to_scheduler() {
    reset_log();
    let stack = Stack::new(64 * 1024).unwrap();
    let sp = init_actor_stack(stack.top(), actor_simple);
    set_actor_sp(sp);
    unsafe { switch_to_actor() };
    assert_eq!(get_log(), 0x1);
}

extern "C-unwind" fn actor_two_steps() {
    log(0x1);
    unsafe { switch_to_scheduler() };
    log(0x2);
    unsafe { switch_to_scheduler() };
}

#[test]
fn actor_yields_and_resumes() {
    reset_log();
    let stack = Stack::new(64 * 1024).unwrap();
    let sp = init_actor_stack(stack.top(), actor_two_steps);
    set_actor_sp(sp);

    unsafe { switch_to_actor() };
    assert_eq!(get_log(), 0x1, "after first resume");

    unsafe { switch_to_actor() };
    assert_eq!(get_log(), 0x1 | 0x2, "after second resume");
}

// Callee-saved registers must survive a yield.

use std::sync::OnceLock;

static REG_BEFORE: OnceLock<[u64; 4]> = OnceLock::new();
static REG_AFTER:  OnceLock<[u64; 4]> = OnceLock::new();

#[cfg(target_arch = "x86_64")]
extern "C-unwind" fn actor_reg_check() {
    unsafe {
        let s0: u64 = 0xAAAA_BBBB_0000_0001;
        let s1: u64 = 0xCCCC_DDDD_0000_0002;
        let s2: u64 = 0xEEEE_FFFF_0000_0003;
        let s3: u64 = 0x1111_2222_0000_0004;

        core::arch::asm!(
            "mov r12, {s0}", "mov r13, {s1}", "mov r14, {s2}", "mov r15, {s3}",
            s0 = in(reg) s0, s1 = in(reg) s1, s2 = in(reg) s2, s3 = in(reg) s3,
            out("r12") _, out("r13") _, out("r14") _, out("r15") _,
        );
        REG_BEFORE.set([s0, s1, s2, s3]).ok();
        switch_to_scheduler();

        let a0: u64; let a1: u64; let a2: u64; let a3: u64;
        core::arch::asm!(
            "mov {a0}, r12", "mov {a1}, r13", "mov {a2}, r14", "mov {a3}, r15",
            a0 = out(reg) a0, a1 = out(reg) a1, a2 = out(reg) a2, a3 = out(reg) a3,
        );
        REG_AFTER.set([a0, a1, a2, a3]).ok();
        switch_to_scheduler();
    }
}

// aarch64 sibling: same intent, AAPCS64 callee-saved integer registers.
// LLVM reserves x19 (and uses x29/x30 as fp/lr), so those can't be named as
// asm operands. We use x23–x26 instead — callee-saved and bindable — moving
// known values into them, yielding, then reading them back, exactly mirroring
// the x86 test's move-into-r12..r15 pattern. The shim saves x23–x26 in its
// `stp x23,x24` / `stp x25,x26` pairs, so a clean round-trip proves they
// survive the switch.
#[cfg(target_arch = "aarch64")]
extern "C-unwind" fn actor_reg_check() {
    unsafe {
        let s0: u64 = 0xAAAA_BBBB_0000_0001;
        let s1: u64 = 0xCCCC_DDDD_0000_0002;
        let s2: u64 = 0xEEEE_FFFF_0000_0003;
        let s3: u64 = 0x1111_2222_0000_0004;

        core::arch::asm!(
            "mov x23, {s0}", "mov x24, {s1}", "mov x25, {s2}", "mov x26, {s3}",
            s0 = in(reg) s0, s1 = in(reg) s1, s2 = in(reg) s2, s3 = in(reg) s3,
            out("x23") _, out("x24") _, out("x25") _, out("x26") _,
        );
        REG_BEFORE.set([s0, s1, s2, s3]).ok();
        switch_to_scheduler();

        let a0: u64; let a1: u64; let a2: u64; let a3: u64;
        core::arch::asm!(
            "mov {a0}, x23", "mov {a1}, x24", "mov {a2}, x25", "mov {a3}, x26",
            a0 = out(reg) a0, a1 = out(reg) a1, a2 = out(reg) a2, a3 = out(reg) a3,
        );
        REG_AFTER.set([a0, a1, a2, a3]).ok();
        switch_to_scheduler();
    }
}

#[test]
fn callee_saved_registers_survive_yield() {
    let stack = Stack::new(64 * 1024).unwrap();
    let sp = init_actor_stack(stack.top(), actor_reg_check);
    set_actor_sp(sp);
    unsafe { switch_to_actor(); switch_to_actor(); }
    assert_eq!(REG_BEFORE.get().copied().unwrap(), REG_AFTER.get().copied().unwrap());
}

// Two actors, independent stacks.

thread_local! {
    static A_VAL: Cell<u64> = const { Cell::new(0) };
    static B_VAL: Cell<u64> = const { Cell::new(0) };
}

extern "C-unwind" fn actor_a() {
    A_VAL.with(|c| c.set(0xAAAA));
    unsafe { switch_to_scheduler() };
    let v = A_VAL.with(|c| c.get());
    A_VAL.with(|c| c.set(if v == 0xAAAA { 0xA00D } else { 0xDEAD }));
    unsafe { switch_to_scheduler() };
}

extern "C-unwind" fn actor_b() {
    B_VAL.with(|c| c.set(0xBBBB));
    unsafe { switch_to_scheduler() };
    let v = B_VAL.with(|c| c.get());
    B_VAL.with(|c| c.set(if v == 0xBBBB { 0xB00D } else { 0xDEAD }));
    unsafe { switch_to_scheduler() };
}

#[test]
fn two_actors_dont_corrupt_each_other() {
    let stack_a = Stack::new(64 * 1024).unwrap();
    let stack_b = Stack::new(64 * 1024).unwrap();

    let sp_a = init_actor_stack(stack_a.top(), actor_a);
    let sp_b = init_actor_stack(stack_b.top(), actor_b);

    set_actor_sp(sp_a); unsafe { switch_to_actor() };
    let sp_a = get_actor_sp();

    set_actor_sp(sp_b); unsafe { switch_to_actor() };
    let sp_b = get_actor_sp();

    set_actor_sp(sp_a); unsafe { switch_to_actor() };
    set_actor_sp(sp_b); unsafe { switch_to_actor() };

    assert_eq!(A_VAL.with(|c| c.get()), 0xA00D);
    assert_eq!(B_VAL.with(|c| c.get()), 0xB00D);
}