section .rdata
    INIT_CAPACITY equ 4096

section .text
extern memcpy
extern memmove
extern memswap
extern oom
extern panic
extern eprint_str
extern allocate

global vec_init
global vec_init_with
global vec_extend
global vec_push
global vec_pop
global vec_drop_last
global vec_swap
global vec_remove
global vec_get
global vec_get_or
global vec_drop
global vec_find
global vec_insert
global vec_insert_many
global vec_insert_sorted
global vec_binary_search_by

global vec_tests


;; Byte vector structure
;; start-structs
;; struct BlobVec {
;; data: *mut u8,
;; len: usize,
;; cap: usize,
;; elem_size: usize,
;; drop: Option<extern "C" fn(*mut u8)>,
;; }
;; end-structs
;; size: 40 bytes
;; align: 8 bytes


;; Initialize a new vector
;; rdi: pointer to Vec struct
;; rsi: item_size
;; rdx: drop function pointer
;; fn vec_init(vec: *mut Vec, item_size: usize, drop: Option<fn(*mut u8)>)
vec_init:
    push rbp
    mov rbp, rsp
    push rdx
    push rsi
    push rdi
    ; (*vec).data = allocate(INIT_CAPACITY);
    mov rdi, INIT_CAPACITY
    call allocate
    pop r15
    mov [r15], rax
    ; (*vec).len = 0;
    mov qword [r15 + 8], 0
    pop rsi
    mov qword [r15 + 24], rsi ; item_size
    ; (*vec).capacity = INIT_CAPACITY / item_size;
    mov rax, INIT_CAPACITY
    xor rdx, rdx
    div rsi
    mov qword [r15 + 16], rax
    pop rdx
    mov qword [r15 + 32], rdx ; drop
    pop rbp
    ret

;; Initialize a new vector
;; rdi: pointer to Vec struct
;; rsi: item_size
;; rdx: drop function pointer
;; rcx: capacity
;; fn vec_init_with(vec: *mut Vec, item_size: usize, drop: Option<fn(*mut u8)>, capacity: usize)
vec_init_with:
    push rbp
    mov rbp, rsp
    sub rsp, 0x20
    mov [rsp], rdx              ; save drop
    mov [rsp + 8], rsi          ; save item_size
    mov [rsp + 16], rcx         ; save capacity
    push rdi
    ; (*vec).data = allocate(INIT_CAPACITY);

    xor rdx, rdx
    mov rax, rcx
    mul rsi                     ; capacity * item_size
    mov rdi, rax
    call allocate
    pop r15
    mov [r15], rax
    ; (*vec).len = 0;
    mov qword [r15 + 8], 0
    mov rax, [rsp + 8]
    mov qword [r15 + 24], rax   ; item_size
    mov rax, [rsp + 16]
    mov qword [r15 + 16], rax   ; capacity
    mov rax, [rsp]
    mov qword [r15 + 32], rax   ; drop
    add rsp, 0x20
    pop rbp
    ret

;; rdi: pointer to Vec struct
;; rsi: index
;; rdx: pointer to default value
;; fn vec_get(vec: *mut Vec, index: usize, default: *mut u8) -> *mut u8
vec_get_or:
    push rbp
    mov rbp, rsp
    ; if (index >= vec.len) panic();
    mov rax, [rdi + 8]    ; len
    cmp rsi, rax
    jge .default
    ; return &mut vec.data[index * vec.item_size];
    mov rax, [rdi + 24]   ; item_size
    mul rsi         ; index * item_size
    mov rsi, [rdi]       ; data
    add rax, rsi         ; data + index * item_size
    pop rbp
    ret
.default:
    mov rax, rdx
    pop rbp
    ret

;; rdi: pointer to Vec struct
;; rsi: index
;; fn vec_get(vec: *mut Vec, index: usize) -> *mut u8
vec_get:
    push rbp
    mov rbp, rsp
    ; if (index >= vec.len) panic();
    mov rax, [rdi + 8]    ; len
    cmp rsi, rax
    jge .panic
    ; return &mut vec.data[index * vec.item_size];
    mov rax, [rdi + 24]   ; item_size
    mul rsi         ; index * item_size
    mov rsi, [rdi]       ; data
    add rax, rsi         ; data + index * item_size
    pop rbp
    ret
.panic:
    call panic
    pop rbp
    ret


;; rdi: pointer to Vec struct
;; rsi: pointer to data to push
;; fn vec_push(vec: *mut Vec, data: &[u8])
vec_push:
    push rbp
    mov rbp, rsp
    ; let vec, ptr, size;
    sub rsp, 24
    mov [rsp], rdi          ; save vec
    mov [rsp + 8], rsi      ; save data ptr
    ; vec_try_grow(vec, size);
    mov rsi, [rdi + 8]
    inc rsi
    call vec_try_grow
    ; memcpy(&vec.data[vec.len], data, size);
    mov rax, [rsp]          ; vec
    mov rcx, [rax]          ; vec.data
    mov rdi, [rax + 8]      ; vec.len
    mov rax, [rax + 24]     ; vec.item_size
    mul rdi                 ; len * item_size
    add rax, rcx            ; data + len * item_size
    mov rdi, rax            ; dest ptr
    mov rsi, [rsp + 8]      ; data ptr
    mov rax, [rsp]          ; vec
    mov rdx, [rax + 24]     ; item_size
    call memcpy
    ; vec.len += 1;
    mov rax, [rsp]          ; vec
    add qword [rax + 8], 1
    add rsp, 24
    pop rbp
    ret


;; rdi: pointer to Vec struct
;; fn vec_pop(vec: *mut Vec)
vec_pop:
    push rbp
    mov rbp, rsp
    ; if vec.len == 0 panic();
    mov rax, [rdi + 8]    ; len
    test rax, rax
    je .underflow
    ; vec.len -= 1;
    sub rax, 1
    mov [rdi + 8], rax
    pop rbp
    ret
    ; }
.underflow:
    call panic
    pop rbp
    ret

;; rdi: pointer to Vec struct
;; fn vec_pop(vec: *mut Vec)
vec_drop_last:
    push rbp
    mov rbp, rsp
    ; if vec.len == 0 panic();
    mov rax, [rdi + 8]    ; len
    test rax, rax
    je .underflow
    ; vec.len -= 1;
    sub rax, 1
    mov [rdi + 8], rax
    ; if let Some(drop) = vec.drop {
    mov r8, [rdi + 32]
    test r8, r8
    je .ret
    ;   drop(&mut vec.data[vec.len..vec.len + vec.item_size]);
    xor rdx, rdx
    mul qword [rdi + 24] ; len * item_size
    mov rdi, [rdi]      ; data
    add rdi, rax        ; data + len * item_size
    call r8
.ret:
    pop rbp
    ret
    ; }
.underflow:
    call panic
    pop rbp
    ret


;; rdi: pointer to Vec struct
;; rsi: index a
;; rdx: index b
;; fn vec_swap(vec: *mut Vec, a: usize, b: usize)
vec_swap:
    push rbp
    mov rbp, rsp
    push rbx
    ; let vec, a, b;
    sub rsp, 24
    mov [rsp], rdi              ; save vec
    ; let max_index = max(a, b);
    mov rbx, rsi
    cmp rbx, rdx
    cmovl rbx, rdx
    ; if (max_index >= vec.len) panic();
    mov rax, [rdi + 8]          ; vec.len
    cmp rbx, rax
    jge .panic

    ; a = a * item_size;
    mov rbx, [rdi + 24]   ; item_size
    mov rax, rdx
    xor rdx, rdx
    mul rbx         ; b * item_size
    mov [rsp + 16], rax
    ; b = b * item_size;
    mov rax, rsi
    xor rdx, rdx
    mul rbx         ; a * item_size
    mov [rsp + 8], rax

    ; // SAFETY: a and b must not overlap
    ; memswap(&mut vec.data.add(a), &mut vec.data.add(b), vec.item_size);
    mov rdx, [rdi + 24]   ; item_size
    mov rdi, [rdi]       ; vec.data
    mov rsi, rdi
    add rdi, [rsp + 16]  ; data + a
    add rsi, [rsp + 8] ; data + b
    call memswap
    add rsp, 24
    pop rbx
    pop rbp
    ret
.panic:
    call panic


;; rdi: pointer to Vec struct
;; rsi: index
;; fn vec_remove(vec: *mut Vec, index: usize)
vec_remove:
    push rbp
    mov rbp, rsp
    sub rsp, 0x10
    mov [rsp], rdi          ; save vec
    mov [rsp + 8], rsi      ; save index

    ; if (index >= vec.len) panic();
    mov rax, [rdi + 8]    ; len
    cmp rsi, rax
    jge .panic
    ; if (vec.len - 1 != index {
    sub rax, 1
    cmp rsi, rax
    je .remove_last
    ;   vec_swap(vec, index, vec.len - 1);
    mov rdx, rax          ; len - 1
    call vec_swap
    ; }
    ;   vec.len -= 1;
.remove_last:
    mov rdi, [rsp]      ; vec
    mov rax, [rdi + 8]    ; len
    sub rax, 1
    mov [rdi + 8], rax
    ; if let Some(drop) = vec.drop {
    mov r8, [rdi + 32]
    test r8, r8
    je .ret
    ;   drop(&mut vec.data[vec.len..vec.len + vec.item_size]);
    mul qword [rdi + 24] ; len * item_size
    mov rdi, [rdi]      ; data
    add rdi, rax        ; data + len * item_size
    call r8
    ; }
.ret:
    add rsp, 0x10
    pop rbp
    ret
.panic:
    call panic

;; essentially a reserve() function
;; rdi: pointer to Vec struct
;; rsi: desired size
;; fn vec_try_grow(vec: *mut Vec, new_size: usize) -> bool
vec_try_grow:
    push rbp
    mov rbp, rsp
    push rbx
    mov rax, [rdi + 16]   ; capacity
    mov rbx, [rdi + 8]    ; len
    cmp rsi, rax
    jge .grow                   ; grow if new_size >= capacity
    pop rbx
    xor rax, rax
    pop rbp
    ret
.grow:
    ; let vec;
    ; let new_capacity;
    ; let new_data;
    sub rsp, 24
    mov [rsp], rdi
    ; new_capacity = max(capacity * 2, len + 1)
    mov rcx, [rdi + 16]   ; capacity
    shl rcx, 1            ; capacity * 2
    cmp rcx, rbx
    cmovl rcx, rbx        ; rcx = new_capacity
    mov [rsp + 8], rcx
    mov rax, rcx
    mul qword [rdi + 24]   ; new_capacity * item_size
    ; new_data = allocate(new_capacity);
    mov rdi, rax
    call allocate
    mov [rsp + 16], rax
    mov rdi, [rsp]
    ; memcpy(new_data, old_data, len * vec.item_size);
    mov rax, [rdi + 8]    ; len
    mul qword [rdi + 24]   ; len * item_size
    mov rdx, rax          ; len
    mov rsi, [rdi]        ; old_data
    mov rdi, [rsp + 16]          ; new_data
    call memcpy
    ; free(old_data) // TODO
    ; update vec
    mov rdi, [rsp]
    mov rax, [rsp + 16]
    mov [rdi], rax        ; data = new_data
    mov rcx, [rsp + 8]
    mov [rdi + 16], rcx   ; capacity = new_capacity
    add rsp, 24
    pop rbx
    xor rax, rax
    pop rbp
    ret

;; rdi: pointer to Vec struct
vec_drop:
    push rbp
    mov rbp, rsp
    push r12
    push r13
    push r14
    push rdi
    ; if let Some(drop) = vec.drop {
    mov r12, [rdi + 32]
    test r12, r12
    je .ret
    ;   for i in 0..vec.len {
    xor r13, r13
    mov r14, [rdi + 8]    ; len
.loop:
    cmp r13, r14
    jge .ret
    mov rdi, [rsp]
    ;     drop(&mut vec.data[i * vec.item_size..(i + 1) * vec.item_size]);
    mov rax, r13
    mul qword [rdi + 24] ; i * item_size
    mov rdi, [rdi]      ; data
    add rdi, rax
    call r12
    inc r13
    jmp .loop
    ;   }
    ; }
    ; free(vec.data);
.ret:
    pop rdi
    pop r14
    pop r13
    pop r12
    pop rbp
    ret


;; rdi: pointer to Vec struct
;; rsi: poiter to object to find
;; rdx: compare function fn(ctx: *const (), a: *const T, b: *const T) -> bool
;; rcx: compare_fn context
vec_find:
    push rbp
    mov rbp, rsp
    push r12
    sub rsp, 0x28
    mov [rsp], rdi          ; save vec
    mov [rsp + 8], rsi      ; save object ptr
    mov [rsp + 0x10], rdx     ; save compare fn
    mov rax, [rdi + 24]   ; item_size
    mov [rsp + 0x18], rax     ; save size
    mov [rsp + 0x20], rcx      ; save compare fn context

    xor r12, r12
.loop:
    mov rcx, [rsp]
    mov rcx, [rcx + 8]
    cmp r12, rcx
    jge .not_found
    ; if compare(&vec.data[i], object) {
    mov rdi, [rsp + 0x20]        ; compare fn context
    mov rsi, [rsp + 0x8]
    mov rax, [rsp + 0x18]       ; size
    xor rdx, rdx
    mul r12                     ; i * item_size
    mov rdx, [rsp]              ; vec
    mov rdx, [rdx]             ; vec.data
    add rdx, rax                ; vec.data + i * item_size
    mov rax, [rsp + 0x10]       ; compare fn
    call rax
    test rax, rax
    jne .found
    ; }
    inc r12
    jmp .loop

.not_found:
    mov rax, -1
    jmp .epilogue
.found:
    mov rax, r12

.epilogue:
    add rsp, 0x28
    pop r12
    pop rbp
    ret


;; rdi: pointer to Vec struct
;; rsi : index
;; rdx: pointer to data to insert
;; fn vec_insert(vec: *mut Vec, index: usize, data: *const T)
vec_insert:
    push rbp
    mov rbp, rsp
    push r12
    push rbx
    sub rsp, 0x20
    mov [rsp], rdi          ; save vec
    mov [rsp + 8], rsi      ; save index
    mov [rsp + 0x10], rdx     ; save data ptr

    mov rsi, [rdi + 8]
    inc rsi
    call vec_try_grow

    mov rbx, [rsp]          ; vec
    mov rdi, [rbx]          ; vec.data
    mov rcx, [rbx + 24]     ; item_size
    mov rdx, [rbx + 8]      ; len
    sub rdx, [rsp + 8]      ; len - index
    mov rax, rdx
    mul rcx                 ; (len - index) * item_size
    mov r12, rax            ; number of bytes to move
    mov rsi, [rbx]          ; vec.data
    mov rax, [rsp + 8]      ; index
    mov rcx, [rbx + 24]     ; item_size
    mul rcx                 ; index * item_size
    add rsi, rax            ; src ptr
    mov rdi, rsi
    add rdi, rcx            ; dst ptr
    mov rdx, r12            ; number of bytes to move
    call memmove

    mov rdi, rsi
    mov rsi, [rsp + 0x10]
    mov rdx, [rbx + 24]     ; item_size
    call memcpy

    mov rax, [rbx + 8]      ; len
    inc rax
    mov [rbx + 8], rax

    add rsp, 32
    pop rbx
    pop r12
    pop rbp
    ret

;; rdi: pointer to Vec struct
;; rsi : index
;; rdx: pointer to data to insert
;; rcx: count
;; define-fn: fn vec_insert_many(vec: *mut BlobVec, index: usize, data: *const u8, count: usize)
vec_insert_many:
    push rbp
    mov rbp, rsp
    push r12
    push rbx
    sub rsp, 0x20
    mov [rsp], rdi              ; save vec
    mov [rsp + 8], rsi          ; save index
    mov [rsp + 0x10], rdx       ; save data ptr
    mov [rsp + 0x18], rcx       ; save count

    mov rsi, [rdi + 8]
    add rsi, rcx                ; len + count
    call vec_try_grow

    mov rbx, [rsp]          ; vec
    mov rdi, [rbx]          ; vec.data
    mov rcx, [rbx + 24]     ; item_size
    mov rdx, [rbx + 8]      ; len
    sub rdx, [rsp + 8]      ; len - index
    mov rax, rdx
    mul rcx                 ; (len - index) * item_size
    mov r12, rax            ; number of bytes to move
    mov rsi, [rbx]          ; vec.data
    mov rax, [rsp + 8]      ; index
    mov rcx, [rbx + 24]     ; item_size
    mul rcx                 ; index * item_size
    add rsi, rax            ; src ptr
    mov rdi, rsi

    mov rax, [rsp + 0x18]     ; count
    mul rcx                 ; count * item_size
    add rdi, rax            ; dst ptr
    mov rdx, r12            ; number of bytes to move
    call memmove

    mov rdi, rsi
    mov rsi, [rsp + 0x10]
    mov rcx, [rsp + 0x18]     ; count
    mov rax, [rbx + 24]     ; item_size
    mul rcx                 ; count * item_size
    mov rdx, rax
    call memcpy

    mov rax, [rbx + 8]      ; len
    add rax, [rsp + 0x18]    ; len += count
    mov [rbx + 8], rax

    add rsp, 32
    pop rbx
    pop r12
    pop rbp
    ret


;; rdi: pointer to Vec struct
;; rsi: pointer to key
;; rdx: compare function fn(ctx: *const (), a: *const T, b: *const T) -> i32
;; rcx: compare_fn context
;;fn vec_binary_search_by(vec: *mut Vec, key: *const T, compare: fn(ctx: *const (), a: *const T, b: *const T) -> i32, ctx: *const ()) -> (usize, bool)
vec_binary_search_by:
    push rbp
    mov rbp, rsp
    sub rsp, 0x40
    mov [rsp], rdi              ; save vec
    mov [rsp + 0x8], rsi        ; save key
    mov [rsp + 0x10], rdx       ; save compare fn
    mov [rsp + 0x18], rcx       ; save compare fn context

    mov rax, [rdi + 8]          ; len
    dec rax                     ; high
    mov qword [rsp + 0x20], 0   ; low
    mov qword [rsp + 0x28], 0   ; mid
    mov [rsp + 0x30], rax       ; high
    mov rax, [rdi + 24]         ; item_size
    mov [rsp + 0x38], rax       ; item_size

.loop:
    mov rax, [rsp + 0x20]      ; low
    mov rcx, [rsp + 0x30]      ; high
    cmp rax, rcx
    jg .not_found
    ; mid = (low + high) / 2
    sub rcx, rax
    shr rcx, 1
    add rax, rcx
    mov [rsp + 0x28], rax       ; mid

    mov rcx, [rsp]              ; vec
    mov rcx, [rcx]              ; vec.data
    mov rdx, [rsp + 0x38]       ; item_size
    mul rdx                     ; mid * item_size
    add rax, rcx                ; vec.data + mid * item_size
    mov rdi, [rsp + 0x18]       ; compare ctx
    mov rsi, [rsp + 0x8]        ; key
    mov rdx, rax
    mov rax, [rsp + 0x10]       ; compare fn
    call rax
    cmp eax, 0
    je .found
    jl .less
    ; greater
    mov rax, [rsp + 0x28]       ; mid
    inc rax
    mov [rsp + 0x28], rax       ; mid
    mov [rsp + 0x20], rax       ; low = mid + 1
    jmp .loop
.less:
    mov rax, [rsp + 0x28]       ; mid
    cmp rax, 0
    je .not_found               ; if mid == 0, not found
    dec rax
    mov [rsp + 0x30], rax       ; high = mid - 1
    jmp .loop
.found:
    mov rax, [rsp + 0x28]
    xor rdx, rdx
    jmp .epilogue
.not_found:
    mov rax, [rsp + 0x28]
    mov rdx, 1

.epilogue:
    add rsp, 0x40
    pop rbp
    ret

;; rdi: pointer to Vec struct
;; rsi: pointer to key
;; rdx: compare function fn(ctx: *const (), a: *const T, b: *const T) -> i32
;; rcx: compare_fn context
;;fn vec_insert_sorted(vec: *mut Vec, key: *const T, compare: fn(ctx: *const (), a: *const T, b: *const T) -> i32, ctx: *const ()) -> (usize, bool)
vec_insert_sorted:
    push rbp
    mov rbp, rsp
    sub rsp, 0x18
    mov [rsp], rdi              ; vec
    mov [rsp + 8], rsi          ; element ptr
    call vec_binary_search_by
    test rdx, rdx
    jnz .insert
    ; element already exists
    jmp .epilogue
.insert:
    mov [rsp + 0x10], rax       ; index
    mov rdi, [rsp]              ; vec
    mov rsi, rax                ; index
    mov rdx, [rsp + 8]          ; element ptr
    call vec_insert
    mov rax, [rsp + 0x10]       ; index
    mov rdx, 1                  ; indicate inserted

.epilogue:
    add rsp, 0x18
    pop rbp
    ret

;; rdi: *Vec
;; rsi: *const u8
;; rdx: number of elements
;; define-fn: fn vec_extend(vec: *mut BlobVec, elements: *const u8, count: usize) -> ()
vec_extend:
    push rbp
    mov rbp, rsp

    ; bytes [24..32]
    ; count [16..24]
    ; elements [8..16]
    ; vec [0..8]
    sub rsp, 32

    mov [rsp], rdi              ; vec
    mov [rsp + 8], rsi          ; elements
    mov [rsp + 16], rdx         ; count

    mov rax, [rdi + 24]         ; item_size
    mul rdx                     ; count * item_size

    mov [rsp + 24], rax         ; bytes

    mov rsi, [rsp + 16]         ; count
    add rsi, [rdi + 8]          ; vec.len + count
    call vec_try_grow

    mov rdi, [rsp]              ; vec
    mov rsi, [rdi + 8]          ; vec.len
    mov rax, [rdi + 24]         ; item_size
    mul rsi                     ; vec.len * item_size
    add rax, [rdi]              ; vec.data + vec.len * item_size
    mov rdi, rax                ; dest
    mov rsi, [rsp + 8]          ; elements
    mov rdx, [rsp + 24]         ; bytes
    call memcpy

    mov rdi, [rsp]              ; vec
    mov rax, [rdi + 8]          ; vec.len
    add rax, [rsp + 16]         ; vec.len + count
    mov [rdi + 8], rax

    add rsp, 32
    pop rbp
    ret