from-scratch/lang/tests/vec.rs

#[repr(C)]
pub struct BlobVec {
    pub data: *mut u8,
    pub len: usize,
    pub cap: usize,
    pub elem_size: usize,
    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 as_slice(&self) -> &[T] {
        assert_eq!(self.vec.elem_size, core::mem::size_of::<T>());
        unsafe { core::slice::from_raw_parts(self.vec.data as *const T, self.vec.len) }
    }

    fn as_slice_mut(&mut self) -> &mut [T] {
        assert_eq!(self.vec.elem_size, core::mem::size_of::<T>());
        unsafe { core::slice::from_raw_parts_mut(self.vec.data as *mut T, self.vec.len) }
    }

    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);
        }
    }

    fn insert(&mut self, value: T, index: usize) {
        if index > self.vec.len {
            return;
        }
        let value = core::mem::ManuallyDrop::new(value);
        unsafe {
            vec_insert(
                &mut self.vec,
                index,
                &raw const value as *const T as *const u8,
            );
        }
    }

    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,
                cmp_trampoline::<T, F>,
                &raw mut cmp as *mut F as *mut (),
            );
            if index == usize::MAX {
                None
            } else {
                Some(index)
            }
        }
    }

    fn binary_search_by<F>(&self, elem: &T, mut cmp: F) -> Result<usize, usize>
    where
        F: FnMut(&T, &T) -> i32,
    {
        extern "C" fn cmp_trampoline<T, F: FnMut(&T, &T) -> i32>(
            f: *const (),
            a: *const u8,
            b: *const u8,
        ) -> i32 {
            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, vacant) = vec_binary_search_by(
                &raw const self.vec as *mut _,
                elem as *const T as *const u8,
                cmp_trampoline::<T, F>,
                &raw mut cmp as *mut F as *mut (),
            );
            if vacant {
                Err(index)
            } else {
                Ok(index)
            }
        }
    }

    fn insert_sorted<F>(&self, elem: T, mut cmp: F) -> Result<usize, usize>
    where
        F: FnMut(&T, &T) -> i32,
    {
        extern "C" fn cmp_trampoline<T, F: FnMut(&T, &T) -> i32>(
            f: *const (),
            a: *const u8,
            b: *const u8,
        ) -> i32 {
            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)
        }

        let mut elem = core::mem::ManuallyDrop::new(elem);

        unsafe {
            let (index, inserted) = vec_insert_sorted(
                &raw const self.vec as *mut _,
                &raw const elem as *const u8,
                cmp_trampoline::<T, F>,
                &raw mut cmp as *mut F as *mut (),
            );
            Ok(index)
        }
    }
}

#[unsafe(no_mangle)]
extern "C" fn panic() -> ! {
    panic!("Called panic from external code.");
}

unsafe impl Send for BlobVec {}
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);
    unsafe fn vec_insert(vec: *mut BlobVec, index: usize, elem: *const u8);
    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,
        cmp: extern "C" fn(*const (), *const u8, *const u8) -> bool,
        cmp_data: *mut (),
    ) -> usize;

    unsafe fn vec_binary_search_by(
        vec: *mut BlobVec,
        elem: *const u8,
        cmp: extern "C" fn(*const (), *const u8, *const u8) -> i32,
        cmp_data: *mut (),
    ) -> (usize, bool);
    unsafe fn vec_insert_sorted(
        vec: *mut BlobVec,
        elem: *const u8,
        cmp: extern "C" fn(*const (), *const u8, *const u8) -> i32,
        cmp_data: *mut (),
    ) -> (usize, bool);
}

fn main() {
    static mut DROPS: usize = 1;
    fn get_drops() -> usize {
        unsafe { (&raw const DROPS).read() }
    }
    unsafe fn update_drops(f: impl FnOnce(&mut usize)) {
        unsafe {
            let drops = &raw mut DROPS;
            f(&mut *drops);
        }
    }

    let mut vec = BlobVec {
        data: core::ptr::null_mut(),
        len: 0,
        cap: 0,
        elem_size: 1,
        drop: None,
    };

    fn as_slice<T>(vec: &BlobVec) -> &[T] {
        assert_eq!(vec.elem_size, core::mem::size_of::<T>());
        unsafe { core::slice::from_raw_parts(vec.data as *const T, vec.len) }
    }

    extern "C" fn drop(ptr: *mut u8) {
        unsafe {
            update_drops(|drops| {
                *drops *= ptr.cast::<u32>().read() as usize;
            });
        }
    }

    unsafe {
        vec_init(&mut vec, 4, Some(drop));
        let mut value: u32 = 2;
        assert_eq!(vec.len, 0);
        vec_push(&mut vec, &value as *const u32 as *const u8);
        assert_eq!(vec.len, 1);
        assert_eq!(as_slice::<u32>(&vec), &[2]);
        let retrieved = *(vec_get(&mut vec, 0) as *mut u32);
        assert_eq!(retrieved, 2);
        assert_eq!(get_drops(), 1);
        vec_drop_last(&mut vec);
        assert_eq!(vec.len, 0);
        assert_eq!(get_drops(), 2);
        value = 3;
        vec_push(&mut vec, &value as *const u32 as *const u8);
        assert_eq!(as_slice::<u32>(&vec), &[3]);
        value = 5;
        vec_push(&mut vec, &value as *const u32 as *const u8);
        assert_eq!(as_slice::<u32>(&vec), &[3, 5]);
        assert_eq!(vec.len, 2);
        vec_drop_last(&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)
    );

    assert_eq!(vec.as_slice(), &[20, 40, 50]);
    vec.insert(30, 1);
    assert_eq!(vec.as_slice(), &[20, 30, 40, 50]);

    let cmp = |a: &u32, b: &u32| match a.cmp(b) {
        core::cmp::Ordering::Less => -1,
        core::cmp::Ordering::Equal => 0,
        core::cmp::Ordering::Greater => 1,
    };
    assert_eq!(vec.binary_search_by(&35, cmp), Err(2));
    assert_eq!(vec.binary_search_by(&25, cmp), Err(1));
    assert_eq!(vec.binary_search_by(&30, cmp), Ok(1));
    assert_eq!(vec.binary_search_by(&5, cmp), Err(0));
    assert_eq!(vec.binary_search_by(&55, cmp), Err(4));

    vec.insert_sorted(35, cmp);
    assert_eq!(vec.as_slice(), &[20, 30, 35, 40, 50]);
}