#include <lib/macros.aarch64_asm.h>

.section .text

// noreturn void enter_in_el1(uint64_t entry, uint64_t sp, uint64_t sctlr,
//                            uint64_t mair, uint64_t tcr, uint64_t ttbr0,
//                            uint64_t ttbr1, uint64_t direct_map_offset)
// Configure EL1 state and jump to kernel. Must be called at EL1.

.global enter_in_el1
enter_in_el1:
    msr spsel, #0
    mov sp, x1

    // Switch to the new page tables

    // Point the EL1t handler to the continuation, such that after we page fault,
    // execution continues and the kernel is entered.
    adrp x8, 1f
    add x8, x8, #:lo12:1f
    add x8, x8, x7
    msr vbar_el1, x8
    isb
    dsb sy
    isb

    // Switch the page table registers
    msr mair_el1, x3
    msr tcr_el1, x4
    msr ttbr0_el1, x5
    msr ttbr1_el1, x6
    msr sctlr_el1, x2
    isb
    dsb sy
    isb

    // Jump to the higher half mapping in case we didn't immediately crash
    br x8

// Alignment required by VBAR register
.align 11
1:
    // Zero out VBAR to avoid confusion
    msr vbar_el1, xzr

    // Disable FP/SIMD/SVE
    msr cpacr_el1, xzr

    // Enter kernel in EL1
    mov x8, #0x3c4
    msr spsr_el1, x8
    msr elr_el1, x0

    mov x0, xzr
    ZERO_REGS_EXCEPT_X0

    eret

// noreturn void enter_in_el2(uint64_t entry, uint64_t sp, uint64_t sctlr,
//                            uint64_t mair, uint64_t tcr, uint64_t ttbr0,
//                            uint64_t ttbr1, uint64_t direct_map_offset)
// Enter kernel at EL2 with VHE. Must be called at EL2.

.global enter_in_el2
enter_in_el2:
    msr spsel, #0
    mov sp, x1

    // Enable E2H if not already set
    mrs x8, hcr_el2
    orr x8, x8, #(1 << 34)
    msr hcr_el2, x8
    isb

    // Switch page tables using VHE-redirected register names.
    // Under VHE, *_el1 writes go to the EL2 register bank.

    // Point the exception handler to the continuation so that if we page fault
    // during the switch, execution continues at the kernel entry.
    adrp x8, 5f
    add x8, x8, #:lo12:5f
    add x8, x8, x7
    msr vbar_el1, x8
    isb
    dsb sy
    isb

    // Switch the page table registers (VHE redirects to EL2)
    msr mair_el1, x3
    msr tcr_el1, x4
    msr ttbr0_el1, x5
    msr ttbr1_el1, x6
    msr sctlr_el1, x2
    isb
    dsb sy
    isb

    // Jump to the higher half mapping in case we didn't immediately crash
    br x8

// Alignment required by VBAR register
.align 11
5:
    // Zero out VBAR to avoid confusion
    msr vbar_el1, xzr

    // Configure EL2 state for VHE

    // Don't trap counters
    mov x8, #3
    msr cnthctl_el2, x8
    msr cntvoff_el2, xzr

    // HCR: E2H + TGE + RW + SWIO
    ldr x8, =0x488000002
    msr hcr_el2, x8

    // Disable FP/SIMD/SVE (VHE CPTR_EL2 layout, CPACR-like)
    msr cptr_el2, xzr
    msr hstr_el2, xzr

    // Enter kernel in EL2
    mov x8, #0x3c8
    msr spsr_el1, x8
    msr elr_el1, x0

    mov x0, xzr
    ZERO_REGS_EXCEPT_X0

    eret

.section .note.GNU-stack,"",%progbits