Better docs and structure

context_switch.S

@@ -1,19 +1,27 @@
 /*
- * context_switch.S – cooperative context switch for x86_64
+ * context_switch.S – Cooperative context switch for x86-64.
  *
  * void context_switch(UINT64 *old_rsp, UINT64 new_rsp);
  *
  *   %rdi = pointer to save current RSP  (old task's PCB field)
  *   %rsi = RSP value to restore         (new task's saved RSP)
  *
- * Saves callee-saved registers and RFLAGS on the current stack,
- * stores RSP, loads the new stack, restores registers, and returns
- * into the new task's execution context.
+ * Saves callee-saved registers (rbp, rbx, r12-r15) and RFLAGS on
+ * the current stack, stores RSP into *old_rsp, loads new_rsp into
+ * RSP, restores registers, and returns into the new task's context.
+ *
+ * This is a standard cooperative switch: only callee-saved state
+ * needs preserving because the C calling convention guarantees that
+ * caller-saved registers are already handled by the compiler.
  */
 
 .text
 .global context_switch
 
+/* ----------------------------------------------------------------
+ *  Entry: save old context, load new context, return
+ * ---------------------------------------------------------------- */
+
 context_switch:
     /* Save callee-saved registers and flags on the current stack */
     pushq   %rbp
@@ -24,10 +32,10 @@ context_switch:
     pushq   %r15
     pushfq
 
-    /* Save current stack pointer into *old_rsp */
+    /* Store current RSP into *old_rsp */
     movq    %rsp, (%rdi)
 
-    /* Load new task's stack pointer */
+    /* Switch to the new task's stack */
     movq    %rsi, %rsp
 
     /* Restore callee-saved registers and flags from the new stack */