; ke_switch.asm - context switch and thread entry stubs (MASM64 / ML64)
;
; Struct offsets:
;   sizeof(OBJECT)  = 16  (type:4, refcount:4, destroy:8)
;   THREAD.SavedContext offset = 88
;   THREAD.TssRsp0  offset     = 104
;   THREAD.Teb      offset     = 112
;   THREAD.FpuContext offset   = 168
;   TSS64.RSP0      offset     = 4
;
THREAD_SavedContext EQU 88
THREAD_TssRsp0      EQU 104
THREAD_Teb          EQU 112
THREAD_FpuContext   EQU 168
TSS_RSP0_OFFSET     EQU 4

EXTERN KeSetCurrentTssRsp0 : PROC

.CODE

; -------------------------------------------------------------------------
; void KeSwitchContext(THREAD* from, THREAD* to)
;   RCX = from, RDX = to (Microsoft x64)
; -------------------------------------------------------------------------
KeSwitchContext PROC
    mov     rax, [rcx + THREAD_FpuContext]
    fxsave64 [rax]

    ; Save non-volatile registers onto the current (outgoing) kernel stack.
    push    rbx
    push    rbp
    push    r12
    push    r13
    push    r14
    push    r15

    ; from->SavedContext = rsp
    mov     [rcx + THREAD_SavedContext], rsp

    ; Update this CPU's TSS.RSP0 before switching to the incoming stack.
    ; The helper hides whether current-CPU storage is still CPU0-only, LAPIC-ID
    ; indexed, or eventually reached through a SWAPGS-backed per-CPU base.
    mov     rbx, rdx
    mov     rcx, [rbx + THREAD_TssRsp0]
    sub     rsp, 40
    call    KeSetCurrentTssRsp0
    add     rsp, 40
    mov     rdx, rbx

    ; rsp = to->SavedContext
    mov     rsp, [rdx + THREAD_SavedContext]

    ; While executing in ring 0, GS_BASE is this CPU's KCPU. Park the
    ; incoming thread's user TEB in IA32_KERNEL_GS_BASE; user-mode entry/exit
    ; swaps it with the kernel GS base using swapgs.
    mov     rax, [rdx + THREAD_Teb] ; Fetch the incoming user GS/TEB pointer
    mov     ecx, 0C0000102h         ; IA32_KERNEL_GS_BASE MSR
    mov     r8, rdx                 ; Save RDX ('to' thread) as WRMSR clobbers EDX
    mov     rdx, rax                ; Copy TEB pointer to RDX
    shr     rdx, 32                 ; EDX = High 32 bits of TEB
                                    ; EAX = Low 32 bits of TEB (upper 32 bits ignored by CPU)
    wrmsr                           ; Write to MSR
    mov     rdx, r8                 ; Restore RDX

    mov     rax, [rdx + THREAD_FpuContext]
    fxrstor64 [rax]

    ; Restore non-volatile registers of the incoming thread.
    pop     r15
    pop     r14
    pop     r13
    pop     r12
    pop     rbp
    pop     rbx
    ret
KeSwitchContext ENDP

; -------------------------------------------------------------------------
; void KeStartInitialThread(THREAD* to)
;   RCX = to
;
; Used by AP scheduler admission. The AP is already in kernel long mode on a
; AP bootstrap stack; this helper installs the idle thread's stack/TSS
; state and returns directly into that thread's saved context. There is no
; outgoing THREAD and no scheduler lock to release.
; -------------------------------------------------------------------------
KeStartInitialThread PROC
    mov     rbx, rcx

    ; Update this CPU's TSS.RSP0 for the incoming thread.
    mov     rcx, [rbx + THREAD_TssRsp0]
    sub     rsp, 40
    call    KeSetCurrentTssRsp0
    add     rsp, 40

    ; Park the incoming thread's user TEB in IA32_KERNEL_GS_BASE. Idle/kernel
    ; threads have a NULL TEB, which is fine.
    mov     rax, [rbx + THREAD_Teb]
    mov     ecx, 0C0000102h
    mov     rdx, rax
    shr     rdx, 32
    wrmsr

    ; Load the thread's saved kernel context and jump into it.
    mov     rsp, [rbx + THREAD_SavedContext]
    mov     rax, [rbx + THREAD_FpuContext]
    fxrstor64 [rax]

    pop     r15
    pop     r14
    pop     r13
    pop     r12
    pop     rbp
    pop     rbx
    ret
KeStartInitialThread ENDP

; -------------------------------------------------------------------------
; KeThreadEntry - landing pad for brand-new threads.
;
; Stack layout when we arrive here (set up by KeCreateThread):
;   [rsp+00] = KCONTEXT (already popped by KeSwitchContext ret)
;   After ret from KeSwitchContext RSP points here:
;     [rsp+00] = user RCX arg
;     [rsp+08] = user RIP
;     [rsp+16] = user CS  (0x1B)
;     [rsp+24] = RFLAGS   (0x202)
;     [rsp+32] = user RSP
;     [rsp+40] = user SS  (0x23)
;
; We load RCX from the stack (first arg per MS x64), then IRETQ.
; -------------------------------------------------------------------------
EXTERN KeSchedulerPostSwitch : PROC
EXTERN KeGetCurrentThread        : PROC
EXTERN KeExitThread           : PROC

KeThreadEntry PROC
    sub     rsp, 32    ; Reserve shadow space
    call    KeSchedulerPostSwitch
    add     rsp, 32
    pop     rcx        ; load user arg
    swapgs             ; GS_BASE: KCPU -> user TEB, KERNEL_GS_BASE: user TEB -> KCPU
    iretq              ; -> ring 3
KeThreadEntry ENDP

; -------------------------------------------------------------------------
; KeKernelThreadEntry - landing pad for kernel-mode threads created by
; KeCreateKernelThread. Never touches ring 3; no IRETQ.
;
; Stack layout when we arrive here (above the KCONTEXT already popped by
; KeSwitchContext's ret):
;   [rsp+00]  fn ptr   (kernel function: void fn(void*))
;   [rsp+08]  arg ptr
;
; We pop both into non-volatile registers so KeSchedulerPostSwitch can
; clobber volatile ones freely, then call fn(arg). On return we call
; KeExitThread so the thread is cleanly destroyed.
; -------------------------------------------------------------------------
KeKernelThreadEntry PROC
    pop     rbx            ; rbx = fn ptr  (non-volatile, survives the call below)
    pop     r12            ; r12 = arg     (non-volatile)

    sub     rsp, 32
    call    KeSchedulerPostSwitch   ; release scheduler spinlock, reap dead threads
    add     rsp, 32

    mov     rcx, r12       ; first arg = arg
    sub     rsp, 32
    call    rbx            ; fn(arg)  -- still ring 0
    add     rsp, 32

    ; Thread function returned -- exit cleanly.
    sub     rsp, 32
    call    KeGetCurrentThread         ; rax = current THREAD*
    add     rsp, 32
    mov     rcx, rax               ; arg1 = thread*
    xor     edx, edx               ; arg2 = exitCode 0
    sub     rsp, 32
    call    KeExitThread
    add     rsp, 32

    ; KeExitThread -> KeYield -> KeSwitchContext never returns here, but
    ; the assembler needs a defined end.
@@: jmp     @B
KeKernelThreadEntry ENDP

END