default rel

section .bss
align 8
free_list: resb 40

section .text
extern memcpy
extern memswap
extern oom
extern panic
extern eprint_error
extern allocate

extern vec_init
extern vec_push
extern vec_get


global bump_init
global bump_alloc

;; rdi: number of bytes to allocate
;; fn mmap(length: usize) -> *mut u8 {
mmap_alloc:
    push rbp
    mov rbp, rsp
    call allocate
    pop rbp
    ret

bump_init:
    push rbp
    mov rbp, rsp

    lea rdi, [rel free_list]
    mov rsi, 16
    mov rdx, 0
    call vec_init

    pop rbp
    ret

;; rdi: min_size
bump_new_block:
    push rbp
    mov rbp, rsp

    cmp rdi, 4096
    ja .mmap
    mov rdi, 4096
.mmap:
    ; next power of 2
    lea rax, [rdi - 1]
    lzcnt rax, rax
    mov rdx, -1
    shrx rdx, rdx, rax
    inc rdx
    mov rdi, rdx
    push rdi
    call mmap_alloc
    pop rdi

    ; Add to free list
    sub rsp, 0x10
    mov [rsp], rax
    mov [rsp + 8], rdi

    lea rdi, [rel free_list]
    mov rsi, rsp
    mov rdx, 0x10
    call vec_push

    add rsp, 0x10
    pop rbp
    ret

;; rdi: number of bytes to allocate
;; rsi: alignment
bump_alloc:
    push rbp
    mov rbp, rsp
    push r12
    push r13
    sub rsp, 0x28
    mov [rsp], rdi      ; size
    mov [rsp + 8], rsi  ; alignment

    lea rdi, [rel free_list]
    mov r12, [rdi + 8]
    xor r13, r13
.alloc_loop:
    cmp r13, r12
    jae .no_block

    lea rdi, [rel free_list]
    mov rsi, r13
    call vec_get
    mov [rsp + 32], rax         ; block entry ptr
    mov rdx, [rax]
    mov [rsp + 16], rdx         ; block_ptr
    mov rax, [rax + 8]          ; block_size
    mov [rsp + 24], rax

    ; check for space
    ; block_ptr.next_multiple_of(alignment) + size <= block_ptr + block_size
    mov rcx, [rsp + 8]    ; alignment
    ; let ptr = block_ptr | alignment
    or rax, rcx                 ; align up
    xor rdx, rdx                ; clear high qword
    ; let rem = ptr % alignment
    div rcx
    ; let diff = alignment - rem
    sub rcx, rdx
    mov rax, [rsp + 16]          ; block+ptr
    ; let aligned_ptr = block_ptr + diff
    add rax, rcx
    ; let aligned_end = aligned_ptr + size
    add rax, [rsp]              ; size
    ; let block_end = block_ptr + block_size
    mov rdx, [rsp + 16]         ; block_ptr
    add rdx, [rsp + 24]         ; block_ptr + block_size
    ; if aligned_end <= block_end
    cmp rax, rdx
    jle .found_space
    ; else try next block
    inc r13
    jmp .alloc_loop
.no_block:
    ; allocate new block
    mov rdi, [rsp]      ; size
    call bump_new_block
    lea r13, [rel free_list]
    mov r13, [r13 + 8]         ; new block index
    mov r12, r13
    dec r13
    jmp .alloc_loop
.found_space:
    sub rdx, rax                ; remaining size in block
    mov r13, [rsp + 32]         ; block entry ptr
    mov [r13], rax
    mov rax, [rsp + 24]         ; block_size
    sub rax, rdx
    mov [r13 + 8], rax
    mov rax, [r13]
    add rsp, 0x28
    pop r13
    pop r12
    pop rbp
    ret