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
extern vec_try_grow_with


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

;; define-fn: fn bump_init()
bump_init:
    push rbp
    mov rbp, rsp

    mov rdi, 0x1000
    call mmap_alloc
    lea rdx, [rel free_list]
    mov [rdx], rax
    mov qword [rdx + 8], 0    ; size 0
    mov rax, 0x1000
    shr rax, 4                  ; capacity in blocks of 16 bytes
    mov qword [rdx + 16], rax   ; capacity
    mov qword [rdx + 24], 16    ; elem_size = 16
    mov qword [rdx + 32], 0     ; drop = None

    pop rbp
    ret

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

    cmp rdi, 4096
    ja .mmap
    mov rdi, 4096
.mmap:

    push rdi

    lea rdi, [rel free_list]
    mov rsi, [rdi + 8]
    mov rax, rsi
    shr rax, 1
    add rax, 1
    add rsi, rax
    mov rdx, mmap_alloc
    call vec_try_grow_with

    pop rdi

    ; 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: num
;; rsi: multiple
next_multiple_of:
    test rsi, rsi
    je .l1
    mov rax, rdi
    or rax, rsi
    shr rax, 32
    je .l2
    mov rax, rdi
    xor edx, edx
    div rsi
    jmp .l3
.l2:
    mov eax, edi
    xor edx, edx
    div esi
.l3:
    sub rsi, rdx
    test rdx, rdx
    cmove rsi, rdx
    add rsi, rdi
    mov rax, rsi
    ret
.l1:
    mov rax, rdi
    ret
    

;; rdi: number of bytes to allocate
;; rsi: alignment
;; define-fn:  fn bump_alloc(size: usize, alignment: usize) -> *mut u8
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 rdi, rdx
    mov rsi, [rsp + 8]          ; alignment
    call next_multiple_of
    ; rax = aligned_ptr

    mov rdx, rax
    add rdx, qword [rsp]        ; ptr + size
    ; rdx = new_end

    mov rcx, [rsp + 16]         ; block_ptr
    add rcx, [rsp + 24]         ; block_size

    cmp rdx, rcx
    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:
    mov r12, [rsp + 32]         ; block entry ptr
    mov rcx, [r12]              ; block_ptr
    mov [r12], rdx              ; update block_ptr

    ; new_end - old_ptr
    sub rdx, rcx

    mov rcx, [r12 + 8]          ; block_size
    sub rcx, rdx
    mov [r12 + 8], rcx

    add rsp, 0x28
    pop r13
    pop r12
    pop rbp
    ret