.option norvc
.option norelax

# ==================================================
# Constants
# ==================================================
.equ PTE_V, 0x001
.equ PTE_R, 0x002
.equ PTE_W, 0x004
.equ PTE_X, 0x008
.equ PTE_A, 0x040
.equ PTE_D, 0x080

.equ SATP_MODE_SV39, (8 << 60)

# ==================================================
# Text
# ==================================================
.section .text
.globl _start

_start:
    # only hart 0
    csrr a0, mhartid
    bnez a0, hang

    # install trap vector (required)
    la   t0, trap_vector
    csrw stvec, t0

    # --------------------------------------------------
    # Build page tables (Sv39 identity map 1 GiB)
    # --------------------------------------------------

    # root_pt[0] -> l1_pt
    la   t0, l1_pt
    srli t0, t0, 12
    slli t0, t0, 10
    ori  t0, t0, PTE_V
    la   t1, root_pt
    sd   t0, 0(t1)

    # l1_pt[0] -> l0_pt
    la   t0, l0_pt
    srli t0, t0, 12
    slli t0, t0, 10
    ori  t0, t0, PTE_V
    la   t1, l1_pt
    sd   t0, 0(t1)

    # l0_pt[0] = 1GiB RWX identity leaf
    li   t0, (PTE_V|PTE_R|PTE_W|PTE_X|PTE_A|PTE_D)
    la   t1, l0_pt
    sd   t0, 0(t1)

    # --------------------------------------------------
    # Enable Sv39
    # --------------------------------------------------
    la   t0, root_pt
    srli t0, t0, 12
    li   t1, SATP_MODE_SV39
    or   t0, t0, t1
    csrw satp, t0
    # sfence.vma zero, zero

    # --------------------------------------------------
    # Init allocator
    # --------------------------------------------------
    la   t0, heap_base
    la   t1, heap_ptr
    sd   t0, 0(t1)

    la   t0, free_list
    sd   zero, 0(t0)

    # --------------------------------------------------
    # Allocate → free → reuse test
    # --------------------------------------------------
    li   a0, 16
    call alloc
    mv   s1, a0

    li   t0, 0x1122334455667788
    sd   t0, 0(s1)

    mv   a0, s1
    call free

    li   a0, 16
    call alloc
    mv   s2, a0

    ld   t1, 0(s2)
    li   t0, 0x1122334455667788
    bne  t1, t0, hang

    # success
    la   t0, tohost
    li   t1, 1
    sd   t1, 0(t0)

hang:
    wfi
    j hang

# ==================================================
# Minimal trap handler (required)
# ==================================================
.align 2
trap_vector:
    csrr t0, scause
    j hang

# ==================================================
# alloc(size)
# ==================================================
.globl alloc
alloc:
    la   t0, free_list
    ld   t1, 0(t0)
    beqz t1, alloc_bump

    ld   t2, 0(t1)
    sd   t2, 0(t0)
    addi a0, t1, 8
    ret

alloc_bump:
    la   t0, heap_ptr
    ld   t1, 0(t0)
    addi t2, a0, 8
    add  t3, t1, t2
    la   t4, heap_end
    bgtu t3, t4, hang
    sd   t3, 0(t0)
    addi a0, t1, 8
    ret

# ==================================================
# free(ptr)
# ==================================================
.globl free
free:
    addi t0, a0, -8
    la   t1, free_list
    ld   t2, 0(t1)
    sd   t2, 0(t0)
    sd   t0, 0(t1)
    ret

# ==================================================
# Data (NO .bss)
# ==================================================
.section .data
.align 3

.globl tohost
.globl fromhost
tohost:   .dword 0
fromhost: .dword 0

.globl heap_ptr
.globl free_list
heap_ptr:  .dword 0
free_list: .dword 0

.align 12
heap_base:
    .zero 4096
heap_end:

.align 12
root_pt: .zero 4096
.align 12
l1_pt:   .zero 4096
.align 12
l0_pt:   .zero 4096