Kernel/Include/Arch/CPU.h

@@ -73,3 +73,16 @@ static inline void CPUEnableMMU(Address l0PhysicalAddress) {
         : "=r"(sctlr) : "r"(sctlr_flags) : "memory"
     );
 }
+
+static inline void CPUSwitchAddressSpace(Address l0Physical) {
+    __asm__ volatile(
+        "dsb ishst\n"             // wait till all previous writes are finished physically
+        "msr ttbr0_el1, %0\n"     // Update TTBR0_EL1 (userspace)
+        "tlbi vmalle1is\n"        // Reset TLB cache
+        "dsb ish\n"               // wait for tlb cache to reset
+        "isb\n"                   // Clear instruction pipeline
+        :: "r" (l0Physical) : "memory"
+    );
+}
+
+#define CPUException(number) __asm__ volatile ("svc %0" :: "i" (number) : "memory")

Kernel/Include/Arch/Exceptions.h

@@ -36,6 +36,8 @@ typedef struct ExceptionsContext {
     UInt64 elr_el1; // pc
     UInt64 spsr_el1; // cpu status
     UInt64 esr_el1; // error reason
+    UInt64 sp_el0; // user's stack
+    UInt64 padding; // align to 16 bytes
 } ExceptionsContext;
 
 typedef enum ExceptionsType {

Kernel/Include/Arch/GIC.h

@@ -25,4 +25,4 @@ enum {
 
 void GICInitialize(Pointer gicdVirtBase, Pointer giccVirtBase);
 void GICEnableInterrupt(UInt32 irqID);
-void GICDispatch(ExceptionsType type);
+Address GICDispatch(ExceptionsContext* frame, ExceptionsType type);

Kernel/Include/Arch/Timer.h

@@ -1,8 +1,11 @@
 #pragma once 
 #include <Types.h>
+#include <Arch/Exceptions.h>
 
 static const UInt64 kTimerFrequency = 1000; // 1ms
 static const UInt8 kTimerIRQ = 27;
 
 void TimerInitialize();
 void TimerReset(UInt64 interval);   
+Address TimerHandler(ExceptionsContext* frame);
+UInt64 TimerGetCounter();

Kernel/Include/OS/Scheduler.h

@@ -0,0 +1,42 @@
+#pragma once
+#include <Types.h>
+
+typedef enum OSTaskState {
+    OSTaskStateDead,
+    OSTaskStateRunning,
+    OSTaskStateReady,
+    OSTaskStateBlocked,
+    OSTaskStateSleeping,
+} OSTaskState;
+
+typedef struct OSProcess {
+    UInt64 id;
+    OSTaskState state;
+    Address* l0Table;
+    Address heapStart;
+    Address heapCurrent;
+    struct OSProcess* parent;
+    ASCII name[32];
+} OSProcess;
+
+typedef struct OSTask {
+    Address stackPointer;
+    struct OSTask* next;
+    UInt32 id;
+    UInt32 sleepTicks;
+    OSTaskState state;
+    Address kernelStackTop;
+    Pointer kernelStackBase;
+    OSProcess* process;
+    Int32 waitingForProcessId;
+} OSTask;
+
+enum {
+    kOSSchedulerTaskStackSize = 16 * 1024,
+    kOSSchedulerExceptionNumber = 0xF0F0
+};
+
+void SchedulerInitialize();
+OSTask* SchedulerSpawn(void(*entryPoint)(), OSProcess* owner, Boolean isUser, Address fixedUserStackAddress);
+Address SchedulerNext(Address stackPointer);
+void SchedulerYield(UInt64 ticks);

Kernel/Source/Arch/Exceptions.c

@@ -4,8 +4,8 @@
 #include <IO/Serial.h>
 #include <OS/Panic.h>
 
-void ExceptionsHandler(ExceptionsContext* frame, ExceptionsType type) {
-    if (type == ExceptionsIRQEl1h || type == ExceptionsIRQEl064) return GICDispatch(type);
+Address ExceptionsHandler(ExceptionsContext* frame, ExceptionsType type) {
+    if (type == ExceptionsIRQEl1h || type == ExceptionsIRQEl064) return GICDispatch(frame, type);
     OSPanicException(frame);
 }
 

Kernel/Source/Arch/GIC.c

@@ -52,14 +52,18 @@ void GICCWriteEOIR(UInt32 irqID) {
     GICC[kGICCEOIR / 4] = irqID;
 }
 
-void GICDispatch(ExceptionsType type) {
-    if (type != ExceptionsIRQEl1h && type != ExceptionsIRQEl064) return;
+Address GICDispatch(ExceptionsContext* frame, ExceptionsType type) {
+    if (type != ExceptionsIRQEl1h && type != ExceptionsIRQEl064) return (Address)frame;
     UInt32 irqID = GICCReadIAR() & 0x3FF;
-    if (irqID == 1023) return; // spurious interrupt
+    if (irqID == 1023) return (Address)frame; // spurious interrupt
+
+
+    Address newStackPointer = (Address)frame;
 
     if (irqID == kTimerIRQ) {
-        TimerReset(kTimerFrequency);
+        newStackPointer = TimerHandler(frame);
     }
 
     GICCWriteEOIR(irqID);
+    return newStackPointer;
 }

Kernel/Source/Arch/Timer.c

@@ -1,5 +1,9 @@
 #include <Arch/Timer.h>
 #include <Arch/GIC.h>
+#include <OS/Scheduler.h>
+#include <Arch/Exceptions.h>
+
+static volatile UInt64 sTimerCounter = 0;
 
 void TimerInitialize() {
     GICEnableInterrupt(kTimerIRQ);
@@ -12,3 +16,13 @@ void TimerReset(UInt64 interval) {
     __asm__ volatile ("msr cntv_tval_el0, %0" :: "r"(frequency /interval));
     __asm__ volatile ("msr cntv_ctl_el0, %0" :: "r"((UInt64)1));
 }
+
+Address TimerHandler(ExceptionsContext* frame) {
+    sTimerCounter++;
+    TimerReset(kTimerFrequency);
+    return SchedulerNext((Address)frame);
+}
+
+UInt64 TimerGetCounter() {
+    return sTimerCounter;
+}

Kernel/Source/Arch/vectors.S

@@ -1,6 +1,6 @@
 .macro ventry type
     .align 7
-    sub sp, sp, #272 // save 272 bytes of stack
+    sub sp, sp, #288 // save 288 bytes of stack
     stp x0, x1, [sp, #0] // move stack 
     mov x1, #\type // move type to x1
     b ExceptionsTrapEntry
@@ -50,12 +50,17 @@ ExceptionsTrapEntry:
     mrs x21, elr_el1
     mrs x22, spsr_el1
     mrs x23, esr_el1
+    mrs x24, sp_el0
 
     stp x30, x21, [sp, #16 * 15]
     stp x22, x23, [sp, #16 * 16]
 
     mov x0, sp
     bl ExceptionsHandler
+    mov sp, x0
+
+    ldp x24, xzr, [sp, #16 * 17]
+    msr sp_el0, x24
 
     ldp x22, x23, [sp, #16 * 16]
     msr spsr_el1, x22
@@ -79,7 +84,7 @@ ExceptionsTrapEntry:
     ldp x2, x3, [sp, #16 * 1]
     ldp x0, x1, [sp, #0]
 
-    add sp, sp, #272
+    add sp, sp, #288
     eret
 
 .global ExceptionsVectorsInit

Kernel/Source/KernelMain.c

@@ -8,6 +8,7 @@
 #include <VM/Heap.h>
 #include <OS/Log.h>
 #include <OS/Panic.h>
+#include <OS/Scheduler.h>
 
 void KernelMain(Bootinfo* bootinfo) {
     OSLog("Kernel started.\n");
@@ -16,6 +17,7 @@ void KernelMain(Bootinfo* bootinfo) {
     }
 
     VMBootMemoryMap bootMap = {0};
+    
     bootMap.reservedCount = 0;
     DTBParse(bootinfo->dtb, &bootMap);
 
@@ -30,5 +32,7 @@ void KernelMain(Bootinfo* bootinfo) {
     TimerInitialize();
     CPUEnableInterrupts();
 
+    SchedulerInitialize();
+
     OSLog("Kernel initialized.\n");
 }

Kernel/Source/OS/Panic.c

@@ -107,7 +107,7 @@ __attribute__((noreturn)) void OSPanicException(ExceptionsContext* frame) {
     OSLog("x22 = 0x%X; x23 = 0x%X\n", frame->x22, frame->x23);
     OSLog("x24 = 0x%X; x25 = 0x%X\n", frame->x24, frame->x25);
     OSLog("x26 = 0x%X; x27 = 0x%X\n", frame->x26, frame->x27);
-    OSLog("\t\tx28 = 0x%X\n", frame->x28);
+    OSLog("\t  x28 = 0x%X\n", frame->x28);
     OSLog("FP  = 0x%X; LR  = 0x%X\n", frame->x29, frame->x30);
 
     DrawPanicFooter();

Kernel/Source/OS/Scheduler.c

@@ -0,0 +1,116 @@
+#include <OS/Scheduler.h>
+#include <OS/Panic.h>
+#include <Arch/Exceptions.h>
+#include <Arch/CPU.h>
+#include <VM/VMM.h>
+#include <VM/Heap.h>
+#include <Lib/Align.h>
+#include <Lib/String.h>
+
+OSTask* gOSSchedulerCurrentTask = nullptr;
+UInt32 gOSSchedulerNextProcessID = 1;
+
+static OSProcess sOSSchedulerKernelProcess;
+
+void SchedulerInitialize() {
+    sOSSchedulerKernelProcess.id = 0;
+    sOSSchedulerKernelProcess.state = OSTaskStateRunning;
+    sOSSchedulerKernelProcess.l0Table = gVMKernelL0Table;
+    StringCopy(sOSSchedulerKernelProcess.name, "kernel");
+
+    OSTask* kernelTask = (OSTask*)HeapAllocate(sizeof(OSTask));
+    if (!kernelTask) OSPanic("Failed to allocate kernel task");
+
+    MemorySet(kernelTask, 0, sizeof(OSTask));
+    kernelTask->id = 0;
+    kernelTask->process = &sOSSchedulerKernelProcess;
+    kernelTask->sleepTicks = 0;
+    kernelTask->next = kernelTask;
+    kernelTask->state = OSTaskStateRunning;
+    kernelTask->waitingForProcessId = -1;
+
+    gOSSchedulerCurrentTask = kernelTask;
+}
+
+OSTask* SchedulerSpawn(void(*entryPoint)(), OSProcess* owner, Boolean isUser, Address fixedUserStackAddress) {
+    OSTask* task = (OSTask*)HeapAllocate(sizeof(OSTask));
+    Pointer stackBase = nullptr;
+    if (!task) goto cleanup;
+
+    if (!owner) owner = &sOSSchedulerKernelProcess;
+    Size stackSize = kOSSchedulerTaskStackSize;
+    stackBase = HeapAllocate(stackSize);
+    if (!stackBase) goto cleanup;
+
+    Address stackPointer = (Address)stackBase + stackSize;
+    stackPointer -= sizeof(ExceptionsContext);
+    stackPointer = AlignDown64(stackPointer, 16);
+    ExceptionsContext* context = (ExceptionsContext*)stackPointer;
+    MemorySet(context, 0, sizeof(ExceptionsContext));
+    context->elr_el1 = (Address)entryPoint;
+
+    if (isUser) {
+        context->spsr_el1 = 0x00000000; // El0t (M[3:0] = 0b0000) allow interrupts
+        context->sp_el0 = fixedUserStackAddress;
+    } else {
+        context->spsr_el1 = 0x00000005; // El1h (M[3:0] = 0b0101) allow interrupts
+    }
+
+    task->stackPointer = stackPointer;
+    task->process = owner;
+    task->id = owner->id;
+    task->sleepTicks = 0;
+    task->kernelStackBase = stackBase;
+    task->kernelStackTop = (Address)stackBase + stackSize;
+    task->state = OSTaskStateRunning;
+    task->waitingForProcessId = -1;
+
+    task->next = gOSSchedulerCurrentTask->next;
+    gOSSchedulerCurrentTask->next = task;
+
+    return task;
+cleanup:
+    if (task) HeapFree(task);
+    if (stackBase) HeapFree(stackBase);
+    return nullptr;
+}
+
+Address SchedulerNext(Address stackPointer) {
+    if (!gOSSchedulerCurrentTask) return stackPointer;
+
+    gOSSchedulerCurrentTask->stackPointer = stackPointer;
+    OSTask* taskIterator = gOSSchedulerCurrentTask->next;
+
+    do {
+        if (taskIterator->sleepTicks > 0) taskIterator->sleepTicks--;
+        taskIterator = taskIterator->next;
+    } while (taskIterator != gOSSchedulerCurrentTask->next);
+    
+    if (gOSSchedulerCurrentTask->sleepTicks > 0) gOSSchedulerCurrentTask->sleepTicks--;
+
+    OSTask* nextTask = gOSSchedulerCurrentTask->next;
+    loop {
+        if (nextTask->state == OSTaskStateSleeping && nextTask->sleepTicks == 0) nextTask->state = OSTaskStateRunning;
+        if (nextTask->state == OSTaskStateRunning) break;
+        nextTask = nextTask->next;
+
+        if (nextTask == gOSSchedulerCurrentTask) {
+            if (gOSSchedulerCurrentTask->state == OSTaskStateRunning) break;
+            OSPanic("No running tasks");
+        }
+    }
+
+    if (nextTask->process->l0Table != gOSSchedulerCurrentTask->process->l0Table) {
+        Address physicalL0Table = VMHHDMToPhys((Address)nextTask->process->l0Table);
+        CPUSwitchAddressSpace(physicalL0Table);
+    }
+    gOSSchedulerCurrentTask = nextTask;
+    
+    return gOSSchedulerCurrentTask->stackPointer;
+}
+
+void SchedulerYield(UInt64 ticks) {
+    gOSSchedulerCurrentTask->sleepTicks = ticks;
+    gOSSchedulerCurrentTask->state = OSTaskStateSleeping;
+    CPUException(kOSSchedulerExceptionNumber);
+}