from-scratch/stages/shared/alloc.asm

;; Allocator
global alloc_pages
global dealloc_pages
global realloc_pages

global heap_alloc
global heap_dealloc
global heap_realloc

extern memcpy

section .bss
    heap resq 1

section .text

oom_msg: db "Out of memory!", 10
oom:
    lea rdi, [abs oom_msg]
    mov esi, 15
    mov rax, 1                  ; syscall: write
    mov rdi, 2                  ; fd: stderr
    syscall
    mov rax, 60                 ; syscall: exit
    mov rdi, -1                 ; status: -1
    syscall

;; 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
    jle oom
    ret

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

;; reallocates memory at $rdi[..$rsi] to a new location of size $rdx.
realloc_pages:
    push rbp
    mov rbp, rsp
    sub rsp, 32
    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, 32
    pop rbp
    ret

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 rbp
    mov rbp, rsp
    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
    pop rbp
    ret
.mmap:
    pop rdi
    pop rsi
    call dealloc_pages
    pop rbp
    ret

;; reallocates a chunk of memory at $rdi of size $rsi and alignment $rdx to a new location of size $rcx
heap_realloc:
    push rbp
    mov rbp, rsp
    push rbx
    push r12
    push r13
    push r14

    mov rbx, rdi
    mov r12, rsi
    mov r13, rdx
    mov r14, rcx

    mov rdi, r14
    mov rsi, r13
    call heap_alloc
    mov r14, rax

    mov rdi, rbx
    mov rsi, rax
    mov rdx, r12
    call memcpy

    mov rdi, rbx
    mov rsi, r12
    mov rdx, r13
    call heap_dealloc

    mov rax, r14
    pop r14
    pop r13
    pop r12
    pop rbx
    pop rbp
    ret