.global __x86_64_syscall_vector

.set REG_RAX, 0 * 8
.set REG_RDI, 1 * 8
.set REG_RSI, 2 * 8
.set REG_RDX, 3 * 8
.set REG_R10, 4 * 8
.set REG_R8,  5 * 8
.set REG_R9,  6 * 8

.set REG_RCX, 7 * 8
.set REG_R11, 8 * 8

.set REG_USER_IP, 9 * 8
.set REG_USER_SP, 10 * 8
.set REG_USER_FLAGS, 11 * 8

.set REG_RBX, 12 * 8
.set REG_RBP, 13 * 8  //
.set REG_R12, 14 * 8  //
.set REG_R13, 15 * 8  // Overwritten by iret
.set REG_R14, 16 * 8  //
.set REG_R15, 17 * 8  //

.set GS_TSS, 0x8
.set GS_TMP, 0x10

// 15 general-purpose registers
// user ip
// user sp
// user flags
.set SYSCALL_STATE_SIZE, 18 * 8

.set REG_IRET_RCX, -(SYSCALL_STATE_SIZE - REG_RCX)
.set REG_IRET_R11, -(SYSCALL_STATE_SIZE - REG_R11)
.set REG_IRET_USER_IP, -(SYSCALL_STATE_SIZE - REG_USER_IP)
.set REG_IRET_USER_SP, -(SYSCALL_STATE_SIZE - REG_USER_SP)
.set REG_IRET_USER_FLAGS, -(SYSCALL_STATE_SIZE - REG_USER_FLAGS)

.macro SYSCALL_SAVE_STATE
    subq $SYSCALL_STATE_SIZE, %rsp

    // Syscall-specific ordering for these
    movq %rax, REG_RAX(%rsp)
    movq %rdi, REG_RDI(%rsp)
    movq %rsi, REG_RSI(%rsp)
    movq %rdx, REG_RDX(%rsp)
    movq %r10, REG_R10(%rsp)
    movq %r8, REG_R8(%rsp)
    movq %r9, REG_R9(%rsp)

    movq %rcx, REG_RCX(%rsp)
    movq %r11, REG_R11(%rsp)

    movq %gs:(GS_TMP), %rax
    movq %rcx, REG_USER_IP(%rsp)
    movq %rax, REG_USER_SP(%rsp)
    movq %r11, REG_USER_FLAGS(%rsp)

    movq %rbx, REG_RBX(%rsp)
    movq %rbp, REG_RBP(%rsp)
    movq %r12, REG_R12(%rsp)
    movq %r13, REG_R13(%rsp)
    movq %r14, REG_R14(%rsp)
    movq %r15, REG_R15(%rsp)
.endm

.macro SYSCALL_RESTORE_STATE
.endm

.section .text
__x86_64_syscall_vector:
    // On entry:
    // %rcx - userspace %rip
    // %r11 - rflags

    // syscall can only be issued from the userspace, so swapgs
    swapgs

    // Store user RSP
    mov %rsp, %gs:(GS_TMP)
    // Load the task's RSP0 from TSS
    mov %gs:(GS_TSS), %rsp
    mov 4(%rsp), %rsp

    SYSCALL_SAVE_STATE

    mov $0x10, %ax
    mov %ax, %ss
    mov %ax, %ds
    mov %ax, %es

    // Save %rsp to %rdi (arg 1) and %rbp
    mov %rsp, %rdi
    mov %rsp, %rbp

    // Force correct stack alignment
    orq $0xF, %rsp
    xorq $0xF, %rsp

    call {syscall_handler}

    // Restore the stack
    mov %rbp, %rsp

.restore_state:
    // Restore non-clobbered state
    movq REG_RAX(%rsp), %rax
    movq REG_RDI(%rsp), %rdi
    movq REG_RSI(%rsp), %rsi
    movq REG_RDX(%rsp), %rdx
    movq REG_R10(%rsp), %r10
    movq REG_R8(%rsp), %r8
    movq REG_R9(%rsp), %r9

    movq REG_RBX(%rsp), %rbx
    movq REG_RBP(%rsp), %rbp
    movq REG_R12(%rsp), %r12
    movq REG_R13(%rsp), %r13
    movq REG_R14(%rsp), %r14
    movq REG_R15(%rsp), %r15

    movq REG_RCX(%rsp), %rcx
    movq REG_USER_IP(%rsp), %r11

    // TODO do I also need to check if rflags changed?
    // If user_ip != rcx (syscall user ip), then use iretq instead
    // Most likely the frame was loaded from a signal entry/return
    cmpq %rcx, %r11
    jne .return_via_iret

    // Still not restored:
    // %rcx (user ip), %r11 (user flags), user sp

.return_via_sysret:
    // Regular syscall return
    movq REG_USER_SP(%rsp), %rcx
    movq %rcx, %gs:(GS_TMP)

    movq REG_USER_IP(%rsp), %rcx
    movq REG_USER_FLAGS(%rsp), %r11

    addq $SYSCALL_STATE_SIZE, %rsp

    // %rcx and %r11 now contain the expected values
    // Restore user RSP
    mov %gs:(GS_TMP), %rsp

    swapgs
    sysretq

.return_via_iret:
    .set IRET_FRAME_SIZE, 5 * 8
    .set IRET_SS, 4 * 8
    .set IRET_RSP, 3 * 8
    .set IRET_RFLAGS, 2 * 8
    .set IRET_CS, 1 * 8
    .set IRET_RIP, 0 * 8

    // Need to restore %rcx, %r11, user ip, user sp, user flags
    // Syscall frame is ordered in a way to prevent iret frame from
    // overwriting any context that was not restored at this moment

    // r15, r14, r13, r12, and rbp will be overwritten, but they're
    // already restored to their registers by now

    // Restore %r11 and only use %rcx
    movq REG_R11(%rsp), %r11

    addq $SYSCALL_STATE_SIZE, %rsp

    subq $IRET_FRAME_SIZE, %rsp

    // SS:RSP
    movq (IRET_FRAME_SIZE + REG_IRET_USER_SP)(%rsp), %rcx
    movq $0x1B, IRET_SS(%rsp)
    movq %rcx, IRET_RSP(%rsp)

    // RFLAGS
    movq (IRET_FRAME_SIZE + REG_IRET_USER_FLAGS)(%rsp), %rcx
    movq %rcx, IRET_RFLAGS(%rsp)

    // CS:RIP
    movq (IRET_FRAME_SIZE + REG_IRET_USER_IP)(%rsp), %rcx
    movq $0x23, IRET_CS(%rsp)
    movq %rcx, IRET_RIP(%rsp)

    // Restore %rcx
    movq (IRET_FRAME_SIZE + REG_IRET_RCX)(%rsp), %rcx

    swapgs
    iretq
