Too late but it kinda works

src/scheduler.rs

@@ -0,0 +1,173 @@
+use core::{
+    arch::{naked_asm, riscv64::wfi},
+    array,
+    cell::LazyCell,
+    time::Duration,
+};
+
+use alloc::string::String;
+use log::info;
+
+use crate::{
+    process::{create_processus, ExecutionContext, Process, ProcessState},
+    riscv::enable_supervisor_interrupt,
+    time,
+};
+
+pub const PROCESSUS_COUNT: usize = 16;
+
+pub static mut ACTIVE_PID: usize = 0;
+pub static mut PROCESS_TABLE: LazyCell<[Process; PROCESSUS_COUNT]> = LazyCell::new(|| {
+    array::from_fn(|_| Process {
+        pid: -1,
+        name: String::new(),
+        state: ProcessState::Dead,
+        wake_time: Duration::new(0, 0),
+        ctx: ExecutionContext {
+            ra: core::ptr::null(),
+            sp: core::ptr::null(),
+            gp: 0,
+            tp: 0,
+            a: [0; _],
+            t: [0; _],
+            s: [0; _],
+            mepc: 0,
+            mstatus: 0,
+        },
+        entry_point: None,
+        stack: [0; _],
+    })
+});
+
+pub extern "C" fn idle() {
+    loop {
+        enable_supervisor_interrupt();
+        unsafe {
+            wfi();
+        }
+    }
+}
+
+pub fn scheduler_init() {
+    info!("scheduler init");
+    for pid in 0..PROCESSUS_COUNT {
+        unsafe {
+            PROCESS_TABLE[pid].state = ProcessState::Dead;
+        }
+    }
+
+    create_processus(idle, "idle");
+    unsafe {
+        PROCESS_TABLE[0].state = ProcessState::Active;
+    }
+}
+
+pub fn scheduler(interrupt_state: ExecutionContext) -> usize {
+    // info!("scheduler");
+    unsafe {
+        let prev_pid = ACTIVE_PID;
+
+        if PROCESS_TABLE[ACTIVE_PID].state == ProcessState::Active {
+            PROCESS_TABLE[ACTIVE_PID].state = ProcessState::Activable;
+        }
+
+        loop {
+            if PROCESS_TABLE[ACTIVE_PID].state == ProcessState::Asleep
+                && time::elapsed_time_since_startup() > PROCESS_TABLE[ACTIVE_PID].wake_time
+            {
+                PROCESS_TABLE[ACTIVE_PID].state = ProcessState::Activable;
+            }
+            ACTIVE_PID = (ACTIVE_PID + 1) % PROCESSUS_COUNT;
+            if PROCESS_TABLE[ACTIVE_PID].state == ProcessState::Activable {
+                break;
+            }
+        }
+        PROCESS_TABLE[ACTIVE_PID].state = ProcessState::Active;
+
+        PROCESS_TABLE[prev_pid].ctx = interrupt_state;
+        // PROCESS_TABLE[prev_pid].ctx.t = interrupt_state.t;
+        // PROCESS_TABLE[prev_pid].ctx.ra = interrupt_state.ra;
+        context_switch(
+            PROCESS_TABLE[ACTIVE_PID].entry_point.unwrap(),
+            &raw mut PROCESS_TABLE[prev_pid].ctx,
+            &raw mut PROCESS_TABLE[ACTIVE_PID].ctx,
+        );
+        prev_pid
+    }
+}
+
+#[unsafe(naked)]
+pub extern "C" fn context_switch(
+    code: extern "C" fn(),
+    current: *mut ExecutionContext,
+    next: *mut ExecutionContext,
+) {
+    naked_asm!(
+        "
+    // sd ra, 0(a1)
+    // sd sp, 8(a1)
+    // sd gp, 16(a1)
+    // sd tp, 24(a1)
+    // sd s0, 152(a1)
+    // sd s1, 160(a1)
+    // sd s2, 168(a1)
+    // sd s3, 176(a1)
+    // sd s4, 184(a1)
+    // sd s5, 192(a1)
+    // sd s6, 200(a1)
+    // sd s7, 208(a1)
+    // sd s8, 216(a1)
+    // sd s9, 224(a1)
+    // sd s10, 232(a1)
+    // sd s11, 240(a1)
+    // csrr t0, sepc
+    // sd t0, 248(a1)
+    // csrr t0, sstatus
+    // sd t0, 256(a1)
+
+    # Load next execution context
+    ld t0, 248(a2)
+    csrw sepc, t0
+    ld t0, 256(a2)
+    csrw sstatus, t0
+
+    ld ra, 0(a2)
+    ld sp, 8(a2)
+    ld gp, 16(a2)
+    ld tp, 24(a2)
+    ld a0, 32(a2)
+    ld a1, 40(a2)
+    // Skip a2 since it used as a pointer
+    ld a3, 56(a2)
+    ld a4, 64(a2)
+    ld a5, 72(a2)
+    ld a6, 80(a2)
+    ld a7, 88(a2)
+    ld t0, 96(a2)
+    ld t1, 104(a2)
+    ld t2, 112(a2)
+    ld t3, 120(a2)
+    ld t4, 128(a2)
+    ld t5, 136(a2)
+    ld t6, 144(a2)
+    ld s0, 152(a2)
+    ld s1, 160(a2)
+    ld s2, 168(a2)
+    ld s3, 176(a2)
+    ld s4, 184(a2)
+    ld s5, 192(a2)
+    ld s6, 200(a2)
+    ld s7, 208(a2)
+    ld s8, 216(a2)
+    ld s9, 224(a2)
+    ld s10, 232(a2)
+    ld s11, 240(a2)
+
+    // Restore a2 at the end
+    ld a2, 48(a2)
+
+    addi sp, sp, 264
+
+    sret"
+    );
+}