add item size to vec struct rather than arg

lang/src/vec.asm

@@ -23,13 +23,21 @@ global vec_free
 ;; data: *mut u8,
 ;; len: usize,
 ;; capacity: usize,
+;; item_size: usize,
+;; drop: Option<fn(*mut u8)>,
 ;; }
+;; size: 40 bytes
+;; align: 8 bytes
 
 
 ;; Initialize a new vector
 ;; rdi: pointer to Vec struct
-;; fn vec_init(vec: *mut Vec)
+;; rsi: item_size
+;; rdx: drop function pointer
+;; fn vec_init(vec: *mut Vec, item_size: usize, drop: Option<fn(*mut u8)>)
 vec_init:
+    push rdx
+    push rsi
     push rdi
     ; (*vec).data = allocate(INIT_CAPACITY);
     mov rdi, INIT_CAPACITY
@@ -38,8 +46,14 @@ vec_init:
     mov [r15], rax
     ; (*vec).len = 0;
     mov qword [r15 + 8], 0
-    ; (*vec).capacity = INIT_CAPACITY;
+    pop rsi
-    mov qword [r15 + 16], INIT_CAPACITY
+    mov qword [r15 + 24], rsi ; item_size
+    ; (*vec).capacity = INIT_CAPACITY / item_size;
+    mov rax, INIT_CAPACITY
+    idiv rax, rsi
+    mov qword [r15 + 16], rax
+    pop rdx
+    mov qword [r15 + 32], rdx ; drop
     ret
 
 ;; rdi: pointer to Vec struct
@@ -47,13 +61,17 @@ vec_init:
 ;; rdx: size of data to push
 ;; fn vec_push(vec: *mut Vec, data: &[u8])
 vec_push:
+    ; if data.len() != vec.item_size panic();
+    mov rax, [rdi + 24]
+    cmp rdx, rax
+    jne .panic
     ; let vec, ptr, size;
     sub rsp, 24
     mov [rsp], rdi          ; save vec
     mov [rsp + 8], rsi      ; save data ptr
     mov [rsp + 16], rdx     ; save data size
     ; vec_try_grow(vec, size);
-    mov rsi, rdx
+    mov rsi, [rdi + 24]   ; item_size
     call vec_try_grow
     ; memcpy(&vec.data[vec.len], data, size);
     mov rax, [rsp]          ; vec
@@ -62,24 +80,39 @@ vec_push:
     mov rsi, [rsp + 8]      ; data ptr
     mov rdx, [rsp + 16]     ; data size
     call memcpy
-    ; vec.len += size;
+    ; vec.len += 1;
     mov rax, [rsp]          ; vec
-    mov rbx, [rsp + 16]     ; size
+    add [rax + 8], 1
-    add [rax + 8], rbx
     add rsp, 24
     ret
+.panic:
+    call panic
+    ret
 
 
 ;; rdi: pointer to Vec struct
-;; rsi: size of data to pop
+;; fn vec_pop(vec: *mut Vec)
-;; fn vec_pop(vec: *mut Vec, size: usize)
 vec_pop:
+    ; if vec.len == 0 panic();
     mov rax, [rdi + 8]    ; len
-    cmp rax, rsi
+    test rax, rax
-    jl .underflow
+    je .underflow
-    sub rax, rsi
+    ; 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]);
+    imul rax, [rdi + 24] ; len * item_size
+    mov rsi, [rdi]      ; data
+    add rsi, rax        ; data + len * item_size
+    mov rdi, rsi
+    call r8
+.ret:
     ret
+    ; }
 .underflow:
     call panic
     ret
@@ -88,19 +121,12 @@ vec_pop:
 ;; rdi: pointer to Vec struct
 ;; rsi: index a
 ;; rdx: index b
-;; rcx: step size
+;; fn vec_swap(vec: *mut Vec, a: usize, b: usize)
-;; fn vec_swap(vec: *mut Vec, a: usize, b: usize, step: usize)
 vec_swap:
-    sub rsp, 32
+    push rbx
+    ; let vec, a, b;
+    sub rsp, 24
     mov [rsp], rdi
-    mov [rsp + 8], rcx
-    ; let a = a * step;
-    imul rsi, rcx
-    mov [rsp + 16], rax
-    ; let b = b * step;
-    imul rdx, rcx
-    mov [rsp + 24], rbx
-
     ; let max_index = max(a, b);
     mov rbx, rsi
     cmp rbx, rdx
@@ -109,74 +135,112 @@ vec_swap:
     mov rax, [rdi + 8]    ; len
     cmp rbx, rax
     jge .panic
-    ; // a and b must not overlap
+
+    ; a = a * item_size;
+    mov rbx, [rdi + 24]   ; item_size
+    imul rdx, rbx         ; b * item_size
+    mov [rsp + 16], rax
+    ; b = b * item_size;
+    imul rsi, 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 rcx, [rdi + 24]   ; item_size
     lea rdi, [rdi]       ; vec.data
     mov rsi, rdi
     add rdi [rsp + 16]  ; data + a
     add rsi, [rsp + 24] ; data + b
-    mov rcx, [rsp + 8]   ; step
     call memswap
-    add rsp, 32
+    add rsp, 24
+    pop rbx
+    ret
+.panic:
+    call panic
     ret
 
 
 ;; rdi: pointer to Vec struct
 ;; rsi: index
-;; rdx: step size
+;; fn vec_remove(vec: *mut Vec, index: usize)
-    ;; fn vec_remove(vec: *mut Vec, index: usize, step: usize)
 vec_remove:
-    ; if (index * step >= vec.len) panic();
+    ; if (index >= vec.len) panic();
-    ; if (vec.len - step == index * step {
+    mov rax, [rdi + 8]    ; len
-    ;   vec.len -= step;
+    cmp rsi, rax
-    ; } else {
+    jge .panic
-    ;   vec_swap(vec, index, vec.len / step - 1, step);
+    ; 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 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]);
+    imul rax, [rdi + 24] ; len * item_size
+    lea rdi, [rdi]      ; data
+    add rdi, rax        ; data + len * item_size
+    call r8
+    ; }
+.ret:
+    ret
 
 ;; rdi: pointer to Vec struct
 ;; rsi: push size
 ;; fn vec_try_grow(vec: *mut Vec, push_size: usize) -> bool
 vec_try_grow:
+    push rbx
     mov rax, [rdi + 16]   ; capacity
     mov rbx, [rdi + 8]    ; len
     add rbx, rsi
     cmp rbx, rax
     jg .grow
+    pop rbx
     xor rax, rax
     ret
 .grow:
     ; let vec;
     ; let new_capacity;
-    sub rsp, 16
+    ; let new_data;
+    sub rsp, 24
     mov [rsp], rdi
-    ; new_capacity = max(capacity * 2, len + push_size)
+    ; 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
+    imul rcx, [rdi + 24]   ; new_capacity * item_size
     ; new_data = allocate(new_capacity);
-    mov rdi, rcx
+    mov rdi, rax
     call allocate
+    mov [rsp + 16], rax
     mov rdi, [rsp]
-    ; memcpy(new_data, old_data, len);
+    ; memcpy(new_data, old_data, len * vec.item_size);
-    mov rbx, [rdi + 8]    ; len
+    mov rax, [rdi + 8]    ; len
+    imul rax, [rdi + 24]   ; len * item_size
+    mov rdx, rax          ; len
     mov rsi, [rdi]        ; old_data
-    mov rdi, rax          ; new_data
+    mov rdi, [rsp + 16]          ; new_data
-    mov rdx, rbx          ; len
     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, 16
+    add rsp, 24
+    pop rbx
     xor rax, rax
     ret
 
-
-