.option norvc
.option norelax

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

_start:
vm_boot:
    # Only hart 0 runs
    csrr    a0, mhartid
    bnez    a0, hang

    # --------------------------------------------------
    # Page tables (Sv39, identity map)
    # --------------------------------------------------
    la      t0, l1_pt
    srli    t0, t0, 12
    slli    t0, t0, 10
    ori     t0, t0, 1
    la      t1, root_pt
    sd      t0, 0(t1)

    la      t0, l0_pt
    srli    t0, t0, 12
    slli    t0, t0, 10
    ori     t0, t0, 1
    la      t1, l1_pt
    sd      t0, 0(t1)

    li      t0, (1<<0)|(1<<1)|(1<<2)|(1<<3)|(1<<6)  # V R W X A
    la      t1, l0_pt
    sd      t0, 0(t1)

    # Enable Sv39
    la      t0, root_pt
    srli    t0, t0, 12
    li      t1, (8 << 60)
    or      t0, t0, t1
    csrw    satp, t0
    sfence.vma zero, zero

    # --------------------------------------------------
    # Init allocator (explicit, no .bss)
    # --------------------------------------------------
    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
    jal     ra, alloc
    mv      s1, a0

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

    mv      a0, s1
    jal     ra, free

    li      a0, 16
    jal     ra, 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

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

    # pop free list
    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)
# a0 = user pointer
# ==================================================
.globl free
free:
    addi    t0, a0, -8
    la      t1, free_list
    ld      t2, 0(t1)
    sd      t2, 0(t0)
    sd      t0, 0(t1)
    ret

# ==================================================
# Exit symbol
# ==================================================
.globl exit
exit:
    j exit

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

# Host interface
.globl tohost
.globl fromhost
tohost:     .dword 0
fromhost:   .dword 0

# Allocator state (explicitly initialized)
.globl heap_ptr
.globl free_list
heap_ptr:   .dword 0
free_list:  .dword 0

# Heap
.align 12
heap_base:
    .zero 4096
heap_end:

# Page tables
.align 12
root_pt: .zero 4096
.align 12
l1_pt:   .zero 4096
.align 12
l0_pt:   .zero 4096