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::{init_actor_stack, 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, 4096).unwrap();
    let sp = init_actor_stack(stack.top(), actor_simple);
    let _ = unsafe { switch_to_actor(sp) };
    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, 4096).unwrap();
    let sp = init_actor_stack(stack.top(), actor_two_steps);

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

    let _ = unsafe { switch_to_actor(sp) };
    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();

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();
    }
}

#[test]
fn callee_saved_registers_survive_yield() {
    let stack = Stack::new(64 * 1024, 4096).unwrap();
    let sp = init_actor_stack(stack.top(), actor_reg_check);
    unsafe {
        let sp = switch_to_actor(sp);
        let _ = switch_to_actor(sp);
    }
    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, 4096).unwrap();
    let stack_b = Stack::new(64 * 1024, 4096).unwrap();

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

    let sp_a = unsafe { switch_to_actor(sp_a) };
    let sp_b = unsafe { switch_to_actor(sp_b) };

    let _ = unsafe { switch_to_actor(sp_a) };
    let _ = unsafe { switch_to_actor(sp_b) };

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