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

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, 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() {
    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, 4);
        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, 4);
        assert_eq!(as_slice::<u32>(&vec), &[3]);
        value = 5;
        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_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)
    );
}