default rel

section .bss
    ifile resb 0x10
    buf resb 0x100
    align 8,db 0
    atoms times 24 resb 8
global env
    env resq 1
    env_tail resq 1
    heap resq 0

section .text

global heap_alloc
global heap_dealloc

panic_abort:
    mov rdi, 1
    mov rax, 60
    syscall

;; rdi: *u8
strlen:
    xor rax, rax
.strlen_loop:
    cmp byte [rdi + rax], 0
    je .strlen_done
    inc rax
    jmp .strlen_loop
.strlen_done:
    ret

;; @param lhs: (rdi, rsi)
;; @param rhs: (rdx, rcx)
;; @return al
strcmp:
    cmp rcx, rsi
    cmovb rsi, rcx              ; if rhs is shorter, use its length for the loop
    xor eax, eax
.strcmp_loop:
    cmp rsi, rax
    jz .strcmp_equal
    movzx ecx, byte [rdx + rax]
    cmp byte [rdi + rax], cl
    lea rax, [rax + 1]
    je .strcmp_loop
    seta al                     ; al = lhs > rhs
    sbb al, 0                   ; al = al - CF
    ret
.strcmp_equal:
    xor eax, eax
    ret

;; rdi: src
;; rsi: dst
;; rdx: len
memcpy:
.loop:
    test rdx, rdx
    jz .done
    mov al, byte [rdi]
    mov byte [rsi], al
    inc rsi
    inc rdi
    dec rdx
    jmp .loop
.done:
    ret

;; Source

;; Allocator

;; allocates $rdi bytes worth of pages via mmap
alloc_pages:
    mov rax, 9                  ; syscall: mmap
    mov rsi, rdi                ; length: rdi
    xor rdi, rdi                ; addr: NULL
    mov rdx, 3                  ; prot: PROT_READ | PROT_WRITE
    mov r10, 34                 ; flags: MAP_PRIVATE | MAP_ANONYMOUS
    mov r8, -1                  ; fd: -1
    xor r9, r9                  ; offset: 0
    syscall
    cmp rax, -1
    jae panic_abort
    ret

dealloc_pages:
    mov rax, 11                 ; syscall: munmap
    mov rdi, rsi                ; addr: rsi
    mov rsi, rdx                ; length: rdx
    syscall
    cmp rax, -1
    jae panic_abort
    ret

;; reallocates memory at $rdi[..$rsi] to a new location of size $rdx.
realloc_pages:
    sub rsp, 24
    mov qword [rsp], rsi
    mov qword [rsp + 8], rdi
    mov rdi, rdx
    call alloc_pages
    mov rsi, rdi
    mov rdi, qword [rsp + 8]
    mov rdx, qword [rsp]
    mov qword [rsp + 16], rax
    call memcpy
    mov rdi, qword [rsp + 8]
    mov rsi, qword [rsp]
    call dealloc_pages
    mov rax, qword [rsp + 16]
    add rsp, 24
    ret

;;
;; `heap` is a pointer to a struct of the form struct { [slab; 9] slabs; }
;; when the heap is empty, `heap` is NULL, and the first allocation will allocate 0x1000 bytes for the heap struct, as well as the first slabs
;; slabs have the following form: struct { u64 tail_end; u64* free; block* first_block; }
;; blocks have the following form: struct { [[u8; SIZE]; (PAGESIZE*4-8)/SIZE] chunks; block* next; }

;; slabs need to keep track of free chunks, so the smallest allocation is 0x8 bytes, and they must keep track of the tail of the last block. each block has to keep track of the next block.
;; the correct slab for a given allocation size is log2(next_power_of_two(size)) - 3 such that the first slab is for allocations of at most 0x8 bytes, the second slab is for 0x10 bytes, then 0x20, 0x40, 0x80, ...
;; each block is 4 pages long so that the 0x800 byte slab doesn't waste half its page for the tail pointer.
;; allocations of size 0x1000 or larger are allocated directly via mmap.

make_slab:
    push rdi
    mov rdi, 0x4000
    call alloc_pages
    mov qword [rax + 0x4000 - 8], 0 ; initialize the tail pointer to NULL
    pop rdi
    mov qword [rdi], 0          ; tail_end = 0
    mov qword [rdi + 8], 0      ; free = NULL
    mov qword [rdi + 16], rax   ; first_block = allocated
    mov qword [rdi + 24], rax   ; last_block = allocated
    mov rax, rdi
    ret

init_heap:
    push r14
    xor r14, r14
    mov rax, qword [rel heap]
    test rax, rax
    jnz .done
    mov rdi, 0x1000
    call alloc_pages
    mov qword [rel heap], rax
.loop:
    cmp r14, 9
    jge .done
    mov rdi, r14
    shl rdi, 5                  ;  idx * 32
    mov rax, qword [rel heap]
    lea rdi, [rax + rdi]        ;  &heap.slabs[idx]
    call make_slab
    inc r14
    jmp .loop
.done:
    pop r14
    ret

;; finds the correct slab for an allocation with size $rdi and align $rsi
slab_bucket:
    xor rax, rax
    dec rdi                     ; if size is a power of two, dec so we can later inc
    dec rsi                     ; ^^
    or rdi, rsi                 ; we only care about the log2, so just gather all the bits
    bsr rsi, rdi                ; log2((size-1) | (align-1))
    sub rsi, 2                  ; +1 for the dec, -3 to collapse the first 3 slabs into one
    cmovae rax, rsi             ; saturating sub
    ret

slab_alloc:
    push rbx
    mov rax, qword [rdi + 8]        ; free
    test rax, rax
    jz .no_free
    mov rdx, qword [rax]            ; next free chunk
    mov qword [rdi + 8], rdx        ; free = next
    pop rbx
    ret
.no_free:
    add esi, 3                  ; undo the -3 from slab_bucket to get the actual log2(size)
    and esi, 63                 ; clamp for safety
    mov edx, 16376              ; 0x4000 - 8
    mov ecx, esi
    shr rdx, cl                 ; 0x4000 - 8 >> log2(size)
    mov rax, qword [rdi]        ; tail_end
    mov rbx, qword [rdi + 24]   ; last_block
    cmp rax, rdx
    jb .alloc_from_block
    push rsi
    push rax
    push rdi
    mov rdi, 0x4000             ; allocate a new block
    call alloc_pages
    mov qword [rax + 0x4000 - 8], 0 ; initialize the tail pointer to NULL
    pop rdi
    mov rsi, qword [rdi + 24]   ; last_block
    mov qword [rsi + 0x4000 - 8], rax
    mov qword [rdi + 24], rax   ; last_block = new block
    mov rbx, rax
    pop rax
    pop rsi
.alloc_from_block:
    inc qword [rdi]             ; tail_end++
    mov ecx, esi
    shl rax, cl
    add rax, rbx
    pop rbx
    ret

    ;; allocate a chunk of memory of size $rdi and alignment $rsi
heap_alloc:
    mov rax, qword [rel heap]
    cmp rax, 0
    je .init
.is_init:
    push rdi
    call slab_bucket
    cmp rax, 9
    jge .mmap
    mov rsi, rax
    shl rax, 5                  ; idx * 32
    mov rdi, qword [rel heap]
    lea rdi, [rdi + rax]        ; &heap.slabs[idx]
    call slab_alloc
    pop rdi
    ret
.init:
    push rdi
    push rsi
    call init_heap
    pop rsi
    pop rdi
    jmp .is_init
.mmap:
    pop rdi
    call alloc_pages
    ret

;; deallocates a chunk of memory at $rdi of size $rsi and alignment $rdx
heap_dealloc:
    push rsi
    push rdi
    mov rdi, rsi
    mov rsi, rdx
    call slab_bucket
    cmp rax, 9
    jge .mmap
    mov rsi, rax
    shl rax, 5                  ; idx * 32
    mov rdi, qword [rel heap]
    lea rdi, [rdi + rax]        ; &heap.slabs[idx]
    mov rax, qword [rdi + 8]    ; free
    pop rsi
    mov qword [rsi], rax
    mov qword [rdi + 8], rsi
    pop rsi
    ret
.mmap:
    pop rdi
    pop rsi
    call dealloc_pages
    ret