more fixes

arch/x86_64/sched.c

@@ -1,4 +1,5 @@
 #include <zinnia/alloc.h>
+#include <zinnia/assert.h>
 #include <zinnia/compiler.h>
 #include <zinnia/errno.h>
 #include <zinnia/percpu.h>
@@ -58,7 +59,7 @@ struct task* arch_sched_switch(struct task* from, struct task* to) {
     struct arch_task_context* from_ctx = &from->context;
     struct arch_task_context* to_ctx = &to->context;
 
-    cpu->tss.rsp0 = to->kernel_stack + KERNEL_STACK_SIZE;
+    cpu->tss.rsp0 = kernel_stack_top(&to->kernel_stack);
 
     if (from_ctx->fpu_region != nullptr) {
         cpu->fpu_save(from_ctx->fpu_region);
@@ -90,15 +91,17 @@ struct task* arch_sched_switch(struct task* from, struct task* to) {
     return perform_switch(&from_ctx->rsp, to_ctx->rsp, from);
 }
 
+// Calls task_entry with the correct parameters.
 [[__naked]]
 void task_entry_thunk() {
     asm volatile(
         "mov rdi, rbx\n"
         "mov rsi, r12\n"
+        "mov rdx, rax\n"
         "jmp %0"
         :
         : "r"(task_entry)
-        : "memory"
+        : "memory", "rdi", "rsi", "rdx"
     );
 }
 
@@ -111,8 +114,10 @@ errno_t arch_task_init(
 ) {
     struct arch_percpu* cpu = percpu_ref(arch_data);
 
+    assert((stack_start & 0xF) == 0, "Task stack is not 16-byte aligned!\n");
+
     // Prepare a dummy stack with an entry point function to return to.
-    struct task_frame* frame = (struct task_frame*)stack_start - 1;
+    struct task_frame* frame = (struct task_frame*)(stack_start - sizeof(uint64_t)) - 1;
     frame->rbx = (uint64_t)entry;
     frame->r12 = arg0;
     frame->rip = (uint64_t)task_entry_thunk;
@@ -133,6 +138,7 @@ errno_t arch_task_init(
         for (size_t i = 0; i < fpu_size; i += arch_mem_page_size())
             pmap_map(&kernel_vma.pmap, fpu_base + i, pages + i, PTE_WRITE | PTE_READ, CACHE_WRITE_BACK);
         context->fpu_region = (void*)fpu_base;
+        context->fpu_phys = pages;
 
         context->ds = asm_read_ds();
         context->es = asm_read_es();
@@ -145,6 +151,21 @@ errno_t arch_task_init(
     return 0;
 }
 
+void arch_task_fini(struct arch_task_context* context) {
+    if (context->fpu_region == nullptr)
+        return;
+
+    const size_t page_size = arch_mem_page_size();
+    const size_t fpu_size = percpu_ref(arch_data)->fpu_size;
+    const uintptr_t fpu_base = (uintptr_t)context->fpu_region;
+    for (size_t i = 0; i < fpu_size; i += page_size)
+        pmap_unmap(&kernel_vma.pmap, fpu_base + i);
+    mem_phys_free(context->fpu_phys, ALIGN_UP(fpu_size, page_size) / page_size);
+    vm_free(fpu_base, fpu_size);
+
+    context->fpu_region = nullptr;
+}
+
 [[__naked]]
 void arch_sched_preempt_disable() {
     asm volatile(