.section .text,"ax"

.equ X19_OFFSET, 0x00
.equ X20_OFFSET, 0x08
.equ X21_OFFSET, 0x10
.equ X22_OFFSET, 0x18
.equ X23_OFFSET, 0x20
.equ X24_OFFSET, 0x28
.equ X25_OFFSET, 0x30
.equ X26_OFFSET, 0x38
.equ X27_OFFSET, 0x40
.equ X28_OFFSET, 0x48
.equ X29_OFFSET, 0x50
.equ X30_OFFSET, 0x58

.equ SP_OFFSET, 0x60

.equ ELR_OFFSET, 0x68
.equ SPSR_OFFSET, 0x70

.equ FPU_REG_OFFSET, 0x80

.equ FPCR_OFFSET, 0x280
.equ FPSR_OFFSET, 0x288





.global __context_switch
__context_switch:
    // x0 -> pointer to current task_context_t struct
    // x1 -> pointer to new task_context_t struct

    // Save currnet task context   str x19, [x0, #X19_OFFSET]
    str x20, [x0, #X20_OFFSET]
    str x21, [x0, #X21_OFFSET]
    str x22, [x0, #X22_OFFSET]
    str x23, [x0, #X23_OFFSET]
    str x24, [x0, #X24_OFFSET]
    str x25, [x0, #X25_OFFSET]
    str x26, [x0, #X26_OFFSET]
    str x27, [x0, #X27_OFFSET]
    str x28, [x0, #X28_OFFSET]
    str x21, [x0, #X21_OFFSET]
    str x22, [x0, #X22_OFFSET]
    str x23, [x0, #X23_OFFSET]
    str x24, [x0, #X24_OFFSET]
    str x25, [x0, #X25_OFFSET]
    str x26, [x0, #X26_OFFSET]
    str x27, [x0, #X27_OFFSET]
    str x28, [x0, #X28_OFFSET]
    str x29, [x0, #X29_OFFSET]
    str x30, [x0, #X30_OFFSET]

    // Store stack pointer
    mov x19, sp
    str x19, [x0, #SP_OFFSET]

    // Store exception link register and saved program status register
    mrs x19, elr_el1
    str x19, [x0, #ELR_OFFSET]

    mrs x19, spsr_el1
    str x19, [x0, #SPSR_OFFSET]

    // Store FPU registers
    stp q0, q1, [x0, #(8 * 32)]
    stp q2, q3, [x0, #(9 * 32)]
    stp q4, q5, [x0, #(10 * 32)]
    stp q6, q7, [x0, #(11 * 32)]
    stp q8, q9, [x0, #(12 * 32)]
    stp q10, q11, [x0, #(13 * 32)]
    stp q12, q13, [x0, #(14 * 32)]
    stp q14, q15, [x0, #(15 * 32)]
    stp q16, q17, [x0, #(16 * 32)]
    stp q18, q19, [x0, #(17 * 32)]
    stp q20, q21, [x0, #(18 * 32)]
    stp q22, q23, [x0, #(19 * 32)]
    stp q24, q25, [x0, #(20 * 32)]
    stp q26, q27, [x0, #(21 * 32)]
    stp q28, q29, [x0, #(22 * 32)]
    stp q30, q31, [x0, #(23 * 32)]

    // Load new task context
    ldr x19, [x1, #FPCR_OFFSET]
    msr fpcr, x19

    ldr x19, [x1, #FPSR_OFFSET]
    msr fpsr, x19

    ldr x19, [x1, #SP_OFFSET]
    mov sp, x19

    ldr x19, [x1, #ELR_OFFSET]
    msr elr_el1, x19

    ldr x19, [x1, #SPSR_OFFSET]
    msr spsr_el1, x19

    ldr x19, [x1, #FPCR_OFFSET]
    msr fpcr, x19

    ldr x19, [x1, #FPSR_OFFSET]
    msr fpsr, x19

    ldr x19, [x1, #X19_OFFSET]
    ldr x20, [x1, #X20_OFFSET]
    ldr x21, [x1, #X21_OFFSET]
    ldr x22, [x1, #X22_OFFSET]
    ldr x23, [x1, #X23_OFFSET]
    ldr x24, [x1, #X24_OFFSET]
    ldr x25, [x1, #X25_OFFSET]
    ldr x26, [x1, #X26_OFFSET]
    ldr x27, [x1, #X27_OFFSET]
    ldr x28, [x1, #X28_OFFSET]
    ldr x29, [x1, #X29_OFFSET]
    ldr x30, [x1, #X30_OFFSET]

    ldp q0, q1, [x1, #(8 * 32)]
    ldp q2, q3, [x1, #(9 * 32)]
    ldp q4, q5, [x1, #(10 * 32)]
    ldp q6, q7, [x1, #(11 * 32)]
    ldp q8, q9, [x1, #(12 * 32)]
    ldp q10, q11, [x1, #(13 * 32)]
    ldp q12, q13, [x1, #(14 * 32)]
    ldp q14, q15, [x1, #(15 * 32)]
    ldp q16, q17, [x1, #(16 * 32)]
    ldp q18, q19, [x1, #(17 * 32)]
    ldp q20, q21, [x1, #(18 * 32)]
    ldp q22, q23, [x1, #(19 * 32)]
    ldp q24, q25, [x1, #(20 * 32)]
    ldp q26, q27, [x1, #(21 * 32)]
    ldp q28, q29, [x1, #(22 * 32)]
    ldp q30, q31, [x1, #(23 * 32)]

    // Re-enable interrupts
    msr DAIFClr, #0b1111
    eret

.global __context_restore
__context_restore:
    ldr x19, [x0, #FPCR_OFFSET]
    msr fpcr, x19

    ldr x19, [x0, #FPSR_OFFSET]
    msr fpsr, x19

    ldr x19, [x0, #SP_OFFSET]
    mov sp, x19

    ldr x19, [x0, #ELR_OFFSET]
    msr elr_el1, x19

    ldr x19, [x0, #SPSR_OFFSET]
    msr spsr_el1, x19

    // TODO: Save/load from proper offsets....
    ldp q0, q1, [x0, #(8 * 32)]
    ldp q2, q3, [x0, #(9 * 32)]
    ldp q4, q5, [x0, #(10 * 32)]
    ldp q6, q7, [x0, #(11 * 32)]
    ldp q8, q9, [x0, #(12 * 32)]
    ldp q10, q11, [x0, #(13 * 32)]
    ldp q12, q13, [x0, #(14 * 32)]
    ldp q14, q15, [x0, #(15 * 32)]
    ldp q16, q17, [x0, #(16 * 32)]
    ldp q18, q19, [x0, #(17 * 32)]
    ldp q20, q21, [x0, #(18 * 32)]
    ldp q22, q23, [x0, #(19 * 32)]
    ldp q24, q25, [x0, #(20 * 32)]
    ldp q26, q27, [x0, #(21 * 32)]
    ldp q28, q29, [x0, #(22 * 32)]
    ldp q30, q31, [x0, #(23 * 32)]

    ldr x19, [x0, #X19_OFFSET]
    ldr x20, [x0, #X20_OFFSET]
    ldr x21, [x0, #X21_OFFSET]
    ldr x22, [x0, #X22_OFFSET]
    ldr x23, [x0, #X23_OFFSET]
    ldr x24, [x0, #X24_OFFSET]
    ldr x25, [x0, #X25_OFFSET]
    ldr x26, [x0, #X26_OFFSET]
    ldr x27, [x0, #X27_OFFSET]
    ldr x28, [x0, #X28_OFFSET]
    ldr x29, [x0, #X29_OFFSET]
    ldr x30, [x0, #X30_OFFSET]

    eret