more testing and fixes for vec.asm

lang/src/vec.asm

@@ -10,12 +10,15 @@ extern eprint_str
 extern allocate
 
 global vec_init
+global vec_init_with
 global vec_push
 global vec_pop
+global vec_drop_last
 global vec_swap
 global vec_remove
 global vec_get
 global vec_drop
+global vec_find
 
 global vec_tests
 
@@ -62,6 +65,41 @@ vec_init:
     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
 ;; fn vec_get(vec: *mut Vec, index: usize) -> *mut u8
@@ -129,6 +167,26 @@ vec_push:
 ;; 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();
@@ -143,10 +201,10 @@ vec_pop:
     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 rsi, [rdi]      ; data
+    mov rdi, [rdi]      ; data
-    add rax, rsi        ; data + len * item_size
+    add rdi, rax        ; data + len * item_size
-    mov rdi, rax
     call r8
 .ret:
     pop rbp
@@ -168,13 +226,13 @@ vec_swap:
     push rbx
     ; let vec, a, b;
     sub rsp, 24
-    mov [rsp], rdi
+    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]    ; len
+    mov rax, [rdi + 8]          ; vec.len
     cmp rbx, rax
     jge .panic
 
@@ -189,11 +247,11 @@ vec_swap:
 
     ; // 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
+    mov rdx, [rdi + 24]   ; item_size
-    lea rdi, [rdi]       ; vec.data
+    mov rdi, [rdi]       ; vec.data
     mov rsi, rdi
     add rdi, [rsp + 16]  ; data + a
-    add rsi, [rsp + 24] ; data + b
+    add rsi, [rsp + 8] ; data + b
     call memswap
     add rsp, 24
     pop rbx
@@ -201,8 +259,6 @@ vec_swap:
     ret
 .panic:
     call panic
-    pop rbx
-    ret
 
 
 ;; rdi: pointer to Vec struct
@@ -211,6 +267,10 @@ vec_swap:
 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
@@ -225,6 +285,7 @@ vec_remove:
     ; }
     ;   vec.len -= 1;
 .remove_last:
+    mov rdi, [rsp]      ; vec
     mov rax, [rdi + 8]    ; len
     sub rax, 1
     mov [rdi + 8], rax
@@ -234,17 +295,16 @@ vec_remove:
     je .ret
     ;   drop(&mut vec.data[vec.len..vec.len + vec.item_size]);
     mul qword [rdi + 24] ; len * item_size
-    lea rdi, [rdi]      ; data
+    mov rdi, [rdi]      ; data
     add rdi, rax        ; data + len * item_size
     call r8
     ; }
 .ret:
+    add rsp, 0x10
     pop rbp
     ret
 .panic:
     call panic
-    pop rbp
-    ret
 
 ;; rdi: pointer to Vec struct
 ;; fn vec_try_grow(vec: *mut Vec, push_size: usize) -> bool
@@ -336,3 +396,61 @@ vec_drop:
     pop r12
     pop rbp
     ret
+
+
+;; rdi: pointer to Vec struct
+;; rsi: poiter to object to find
+;; rdx: size of object to find
+;; rcx: compare function fn(ctx: *const (), a: *const T, b: *const T) -> bool
+;; r8: 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], rcx     ; save compare fn
+    mov [rsp + 0x18], rdx     ; save size
+    mov [rsp + 0x20], r8      ; save compare fn context
+
+    mov rcx, [rdi + 24]   ; item_size
+    cmp rdx, rcx
+    jne .panic
+
+    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
+.panic:
+    call panic

lang/tests/vec.rs

@@ -7,6 +7,133 @@ pub struct BlobVec {
     pub drop: Option<extern "C" fn(*mut u8)>,
 }
 
+struct VecT<T> {
+    vec: BlobVec,
+    _marker: core::marker::PhantomData<T>,
+}
+
+impl<T> VecT<T> {
+    fn new() -> Self {
+        Self::new_with(32)
+    }
+
+    fn new_with(capacity: usize) -> Self {
+        let mut vec = BlobVec {
+            data: core::ptr::null_mut(),
+            len: 0,
+            cap: 0,
+            elem_size: 0,
+            drop: None,
+        };
+
+        extern "C" fn drop_fn<T>(ptr: *mut u8) {
+            unsafe {
+                core::ptr::drop_in_place::<T>(ptr as *mut T);
+            }
+        }
+
+        unsafe {
+            vec_init_with(
+                &mut vec,
+                core::mem::size_of::<T>(),
+                Some(drop_fn::<T>),
+                capacity,
+            );
+        }
+
+        Self {
+            vec,
+            _marker: core::marker::PhantomData,
+        }
+    }
+
+    fn push(&mut self, value: T) {
+        let value = core::mem::ManuallyDrop::new(value);
+        unsafe {
+            vec_push(
+                &mut self.vec,
+                &raw const value as *const T as *const u8,
+                core::mem::size_of::<T>(),
+            );
+        }
+    }
+
+    fn pop(&mut self) -> Option<T> {
+        if self.vec.len == 0 {
+            return None;
+        }
+        unsafe {
+            let ptr = vec_get(&mut self.vec, self.vec.len - 1) as *mut T;
+            let value = ptr.read();
+            vec_pop(&mut self.vec);
+            Some(value)
+        }
+    }
+
+    fn get(&self, index: usize) -> Option<&T> {
+        if index >= self.vec.len {
+            return None;
+        }
+        unsafe {
+            let ptr = vec_get(&raw const self.vec as *mut _, index) as *mut T;
+            Some(&*ptr)
+        }
+    }
+    fn get_mut(&mut self, index: usize) -> Option<&mut T> {
+        if index >= self.vec.len {
+            return None;
+        }
+        unsafe {
+            let ptr = vec_get(&raw mut self.vec, index) as *mut T;
+            Some(&mut *ptr)
+        }
+    }
+
+    fn remove(&mut self, index: usize) {
+        if index >= self.vec.len {
+            return;
+        }
+        unsafe {
+            vec_remove(&mut self.vec, index);
+        }
+    }
+
+    fn len(&self) -> usize {
+        self.vec.len
+    }
+
+    fn position<F>(&self, elem: &T, mut cmp: F) -> Option<usize>
+    where
+        F: FnMut(&T, &T) -> bool,
+    {
+        extern "C" fn cmp_trampoline<T, F: FnMut(&T, &T) -> bool>(
+            f: *const (),
+            a: *const u8,
+            b: *const u8,
+        ) -> bool {
+            let f = unsafe { &mut *(f as *mut F) };
+            let a = unsafe { &*(a as *const T) };
+            let b = unsafe { &*(b as *const T) };
+            f(a, b)
+        }
+
+        unsafe {
+            let index = vec_find(
+                &raw const self.vec as *mut _,
+                elem as *const T as *const u8,
+                core::mem::size_of::<T>(),
+                cmp_trampoline::<T, F>,
+                &raw mut cmp as *mut F as *mut (),
+            );
+            if index == usize::MAX {
+                None
+            } else {
+                Some(index)
+            }
+        }
+    }
+}
+
 #[unsafe(no_mangle)]
 extern "C" fn panic() -> ! {
     panic!("Called panic from external code.");
@@ -17,14 +144,29 @@ unsafe impl Sync for BlobVec {}
 
 unsafe extern "C" {
     unsafe fn vec_init(vec: *mut BlobVec, elem_size: usize, drop: Option<extern "C" fn(*mut u8)>);
+    unsafe fn vec_init_with(
+        vec: *mut BlobVec,
+        elem_size: usize,
+        drop: Option<extern "C" fn(*mut u8)>,
+        cap: usize,
+    );
     unsafe fn vec_push(vec: *mut BlobVec, elem: *const u8, size: usize);
     unsafe fn vec_pop(vec: *mut BlobVec);
+    unsafe fn vec_drop_last(vec: *mut BlobVec);
     unsafe fn vec_get(vec: *mut BlobVec, index: usize) -> *mut u8;
 
     #[allow(dead_code)]
     unsafe fn vec_remove(vec: *mut BlobVec, index: usize);
     #[allow(dead_code)]
     unsafe fn vec_drop(vec: *mut BlobVec);
+
+    unsafe fn vec_find(
+        vec: *mut BlobVec,
+        elem: *const u8,
+        size: usize,
+        cmp: extern "C" fn(*const (), *const u8, *const u8) -> bool,
+        cmp_data: *mut (),
+    ) -> usize;
 }
 
 fn main() {
@@ -70,7 +212,7 @@ fn main() {
         let retrieved = *(vec_get(&mut vec, 0) as *mut u32);
         assert_eq!(retrieved, 2);
         assert_eq!(get_drops(), 1);
-        vec_pop(&mut vec);
+        vec_drop_last(&mut vec);
         assert_eq!(vec.len, 0);
         assert_eq!(get_drops(), 2);
         value = 3;
@@ -80,9 +222,34 @@ fn main() {
         vec_push(&mut vec, &value as *const u32 as *const u8, 4);
         assert_eq!(as_slice::<u32>(&vec), &[3, 5]);
         assert_eq!(vec.len, 2);
-        vec_pop(&mut vec);
+        vec_drop_last(&mut vec);
-        vec_pop(&mut vec);
+        vec_drop_last(&mut vec);
         assert_eq!(get_drops(), 2 * 3 * 5);
         eprintln!("Push/pop test passed\n");
     }
+
+    let mut vec = VecT::<u32>::new_with(100);
+    assert_eq!(vec.len(), 0);
+    vec.push(10);
+    vec.push(20);
+    vec.push(30);
+    assert_eq!(vec.len(), 3);
+    assert_eq!(vec.get(0), Some(&10));
+    assert_eq!(vec.get(1), Some(&20));
+    assert_eq!(vec.get(2), Some(&30));
+    assert_eq!(vec.pop(), Some(30));
+    assert_eq!(vec.len(), 2);
+    vec.remove(0);
+    assert_eq!(vec.len(), 1);
+    assert_eq!(vec.get(0), Some(&20));
+    vec.push(40);
+    vec.push(50);
+
+    assert_eq!(
+        vec.position(&40, |a, b| {
+            eprintln!("Comparing {} and {}", a, b);
+            a == b
+        }),
+        Some(1)
+    );
 }