Sylva/kernel/scheduler/scheduler.cpp

#include <scheduler.h>
#include <memory/heap.h>
#include <memory/pmm.h>
#include <common.h>
#include <serial.h>

// Assembly: context_switch(UINT64* old_rsp, UINT64 new_rsp)
extern "C" void context_switch(UINT64* old_rsp, UINT64 new_rsp);

static task_t  g_tasks[TASK_MAX];
static UINT32  g_task_count = 0;
static task_t* g_current = NULL;
static task_t* g_task_list = NULL;  // circular linked list head

// Trampoline: first thing a new task runs after context_switch.
// The entry function pointer is stored in the task's name field
// (we repurpose a slot — actually we store it in a simple global array).
static void (*g_task_entries[TASK_MAX])(void);

extern "C" void task_entry_trampoline() {
    task_t* cur = scheduler_current();
    if (cur && g_task_entries[cur->id]) {
        g_task_entries[cur->id]();   // call the user function
    }
    task_exit();  // clean up when done
}

task_t* task_create(const char* name, void (*entry)(void)) {
    if (g_task_count >= TASK_MAX) {
        serial_write("SCHED: task limit reached\n");
        return NULL;
    }

    UINT32 id = g_task_count++;
    task_t* task = &g_tasks[id];

    // Store entry function for the trampoline
    g_task_entries[id] = entry;

    // Allocate kernel stack
    UINTN stack_pages = TASK_STACK_SIZE / PAGE_SIZE;
    void* stack = pmm_alloc_pages(stack_pages);
    if (!stack) {
        serial_write("SCHED: stack alloc failed for task ");
        serial_write(name);
        serial_write("\n");
        return NULL;
    }

    // Fill task struct
    task->id = id;
    task->state = TASK_STATE_READY;
    task->stack_base = stack;

    // Copy name
    const char* s = name;
    char* d = task->name;
    for (int i = 0; i < TASK_NAME_LEN - 1 && *s; i++) {
        *d++ = *s++;
    }
    *d = '\0';

    // Set up initial stack for first context_switch into this task.
    // Stack grows downward. context_switch will pop 6 regs then ret.
    //
    // Layout (high addr -> low addr):
    //   [stack + TASK_STACK_SIZE]     <- top
    //   return addr = task_entry_trampoline  (ret goes here)
    //   rbx = 0
    //   rbp = 0
    //   r12 = 0
    //   r13 = 0
    //   r14 = 0
    //   r15 = 0                       <- RSP points here initially
    //
    UINT64* sp = (UINT64*)((UINT8*)stack + TASK_STACK_SIZE);

    // Push return address (task_entry_trampoline)
    *--sp = (UINT64)task_entry_trampoline;

    // Push callee-saved registers (all zero)
    *--sp = 0;  // rbx
    *--sp = 0;  // rbp
    *--sp = 0;  // r12
    *--sp = 0;  // r13
    *--sp = 0;  // r14
    *--sp = 0;  // r15

    task->rsp = (UINT64)sp;

    // Insert into circular linked list
    if (g_task_list == NULL) {
        task->next = task;  // points to itself (single element circle)
        g_task_list = task;
    } else {
        // Insert after current tail (g_task_list is the "last" in circle)
        task->next = g_task_list->next;
        g_task_list->next = task;
        g_task_list = task;  // new tail
    }

    serial_write("SCHED: created task '");
    serial_write(task->name);
    serial_write("' id=");
    serial_write_hex(id);
    serial_write("\n");

    return task;
}

void yield(void) {
    if (g_current == NULL || g_task_list == NULL) return;

    task_t* cur = g_current;
    task_t* next = cur->next;

    // Skip terminated tasks
    while (next->state == TASK_STATE_TERMINATED && next != cur) {
        next = next->next;
    }
    if (next->state == TASK_STATE_TERMINATED) return;  // all terminated

    if (next == cur) return;  // only one task, nothing to do

    if (cur->state != TASK_STATE_TERMINATED)
        cur->state = TASK_STATE_READY;
    next->state = TASK_STATE_RUNNING;
    g_current = next;

    context_switch(&cur->rsp, next->rsp);
}

void scheduler_run(void) {
    if (g_task_list == NULL) {
        serial_write("SCHED: no tasks to run\n");
        return;
    }

    // Find first READY task
    task_t* start = g_task_list->next;  // head of circle
    task_t* t = start;
    do {
        if (t->state == TASK_STATE_READY) {
            break;
        }
        t = t->next;
    } while (t != start);

    if (t->state != TASK_STATE_READY) {
        serial_write("SCHED: no READY tasks\n");
        return;
    }

    g_current = t;
    t->state = TASK_STATE_RUNNING;

    serial_write("SCHED: starting task '");
    serial_write(t->name);
    serial_write("'\n");

    // First context switch — no old RSP to save (we're still in scheduler_run)
    // Just switch to the task's stack directly.
    // We need a dummy old_rsp to satisfy the API, but we never return here.
    UINT64 dummy_rsp;
    context_switch(&dummy_rsp, t->rsp);

    // We only return here when ALL tasks are terminated
    serial_write("SCHED: all tasks finished\n");
    while (1) ASM ("hlt");
}

void task_exit(void) {
    if (g_current == NULL) return;

    serial_write("SCHED: task '");
    serial_write(g_current->name);
    serial_write("' exited\n");

    g_current->state = TASK_STATE_TERMINATED;

    // Yield to next task — we won't come back
    yield();

    // Should never reach here
    while (1) ASM ("hlt");
}

task_t* scheduler_current(void) {
    return g_current;
}