lib.rs

#![feature(allocator_api)]
#![feature(strict_provenance)]

use std::{
    alloc::{Allocator, Global, Layout},
    marker::PhantomData,
    mem::ManuallyDrop,
    ops::{Index, IndexMut},
    ptr::NonNull,
};

// use sptr::Strict;

/// A size-optimized Vec. Has a max length of 256 items.
pub struct Vec<T, A: Allocator = Global> {
    /*
        The pointer's bits are laid out as follows:

        bits 0-48: address
        bits 48-56: length
        bits 56-64: capacity
    */
    ptr: *mut T,
    allocator: A,
}

const MASK_LENGTH: usize = 0x00FF_0000_0000_0000;

impl<T> Vec<T, Global> {
    pub const fn new() -> Self {
        Vec {
            ptr: std::ptr::null_mut(),
            allocator: Global,
        }
    }

    pub unsafe fn from_raw_parts(ptr: *mut T, length: usize, capacity: usize) -> Self {
        Vec {
            ptr: tag(ptr, length, capacity),
            allocator: Global,
        }
    }
}

impl<T, A: Allocator> Vec<T, A> {
    pub const fn new_in(allocator: A) -> Self {
        Vec {
            ptr: std::ptr::null_mut(),
            allocator,
        }
    }

    pub unsafe fn from_raw_parts_in(
        ptr: *mut T,
        length: usize,
        capacity: usize,
        allocator: A,
    ) -> Self {
        Vec {
            ptr: tag(ptr, length, capacity),
            allocator,
        }
    }

    pub fn capacity(&self) -> usize {
        self.ptr.addr() >> 56
    }

    pub fn length(&self) -> usize {
        (self.ptr.addr() >> 48) & 0xff
    }

    fn ptr_untagged(&self) -> *mut T {
        // sign extends the pointer
        self.ptr
            .map_addr(|addr| ((addr as isize) << 16 >> 16) as usize)
    }

    pub fn get(&self, idx: usize) -> Option<&T> {
        if idx >= self.length() {
            return None;
        }

        let reference = unsafe { &*self.ptr_untagged().add(idx) };

        Some(reference)
    }

    pub fn get_mut(&mut self, idx: usize) -> Option<&mut T> {
        if idx >= self.length() {
            return None;
        }

        let reference = unsafe { &mut *self.ptr_untagged().add(idx) };

        Some(reference)
    }

    pub fn reserve(&mut self, additional: usize) {
        let length = self.length();
        let old_capacity = self.capacity();
        let capacity = additional + length;
        if capacity >= 256 {
            panic!("cursed_vec::Vec has a max length of 256");
        }

        if capacity > old_capacity {
            let cap = (2 * old_capacity).clamp(capacity.max(4), 256);

            let new_layout = Layout::array::<T>(cap).unwrap();

            let ptr = if let Some(ptr) = NonNull::new(self.ptr_untagged()) {
                let old_layout = Layout::array::<T>(old_capacity).unwrap();
                unsafe {
                    self.allocator
                        .grow(ptr.cast(), old_layout, new_layout)
                        .unwrap()
                }
            } else {
                self.allocator.allocate(new_layout).unwrap()
            };

            self.ptr = tag(ptr.as_ptr() as *mut T, cap, length);
        }
    }

    pub fn push(&mut self, val: T) {
        let idx = self.length();
        if idx >= self.capacity() {
            self.reserve(1);
        }

        unsafe {
            self.ptr_untagged().add(idx).write(val);
        }

        self.ptr = self.ptr.map_addr(|addr| addr + (1 << 48));
    }

    pub fn allocator(&self) -> &A {
        &self.allocator
    }

    pub fn as_mut_ptr(&mut self) -> *mut T {
        self.ptr_untagged()
    }

    pub fn as_ptr(&self) -> *const T {
        self.ptr_untagged()
    }

    pub fn as_mut_slice(&mut self) -> &mut [T] {
        unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.length()) }
    }

    pub fn as_slice(&self) -> &[T] {
        unsafe { std::slice::from_raw_parts(self.as_ptr(), self.length()) }
    }

    pub fn clear(&mut self) {
        for i in 0..self.length() {
            unsafe {
                self.ptr_untagged().add(i).drop_in_place();
            }
        }

        self.ptr = self.ptr.map_addr(|addr| addr & !MASK_LENGTH);
    }

    pub fn into_raw_parts(self) -> (*mut T, usize, usize) {
        let this = ManuallyDrop::new(self);

        (this.ptr_untagged(), this.length(), this.capacity())
    }

    pub fn into_raw_parts_with_alloc(self) -> (*mut T, usize, usize, A) {
        let this = ManuallyDrop::new(self);

        (
            this.ptr_untagged(),
            this.length(),
            this.capacity(),
            unsafe { std::ptr::addr_of!(this.allocator).read() },
        )
    }

    pub fn iter(&self) -> Iter<'_, T> {
        Iter {
            ptr: self.ptr_untagged(),
            index: 0,
            length: self.length(),
            lifetime: PhantomData,
        }
    }

    pub fn iter_mut(&mut self) -> IterMut<'_, T> {
        IterMut {
            ptr: self.ptr_untagged(),
            index: 0,
            length: self.length(),
            lifetime: PhantomData,
        }
    }
}

pub struct Iter<'a, T> {
    ptr: *const T,
    index: usize,
    length: usize,
    lifetime: PhantomData<&'a Vec<T>>,
}

pub struct IterMut<'a, T> {
    ptr: *mut T,
    index: usize,
    length: usize,
    lifetime: PhantomData<&'a mut Vec<T>>,
}

impl<'a, T> Iterator for Iter<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        if self.index == self.length {
            None
        } else {
            let item = unsafe { &*self.ptr.add(self.index) };
            self.index += 1;
            Some(item)
        }
    }
}

impl<'a, T> Iterator for IterMut<'a, T> {
    type Item = &'a mut T;

    fn next(&mut self) -> Option<Self::Item> {
        if self.index == self.length {
            None
        } else {
            let item = unsafe { &mut *self.ptr.add(self.index) };
            self.index += 1;
            Some(item)
        }
    }
}

impl<T> IntoIterator for Vec<T> {
    type Item = T;

    type IntoIter = IntoIter<T>;

    fn into_iter(self) -> Self::IntoIter {
        IntoIter {
            inner: ManuallyDrop::new(self),
            index: 0,
        }
    }
}

pub struct IntoIter<T> {
    inner: ManuallyDrop<Vec<T>>,
    index: usize,
}

impl<T> Iterator for IntoIter<T> {
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        if self.index == self.inner.length() {
            None
        } else {
            let item = unsafe { self.inner.ptr_untagged().add(self.index).read() };
            self.index += 1;
            Some(item)
        }
    }
}

impl<T> Drop for IntoIter<T> {
    fn drop(&mut self) {
        let ptr_untagged = self.inner.ptr_untagged();

        if let Some(ptr) = NonNull::new(ptr_untagged) {
            for i in self.index..self.inner.length() {
                unsafe {
                    ptr_untagged.add(i).drop_in_place();
                }
            }
            // alternatively
            // while let Some(_) = self.next() { /* nothing to see here */ }

            let ptr = ptr.cast();
            let capacity = self.inner.capacity();
            let layout = Layout::array::<T>(capacity).unwrap();

            unsafe {
                self.inner.allocator.deallocate(ptr, layout);
            }
        }
    }
}

impl<T, A: Allocator> Index<usize> for Vec<T, A> {
    type Output = T;

    fn index(&self, index: usize) -> &Self::Output {
        self.get(index).expect("Out of bounds")
    }
}

impl<T, A: Allocator> IndexMut<usize> for Vec<T, A> {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        self.get_mut(index).expect("Out of bounds")
    }
}

fn tag<T>(ptr: *mut T, cap: usize, length: usize) -> *mut T {
    ptr.map_addr(|addr| addr | (cap << 56) | (length << 48))
}

impl<T, A: Allocator> Drop for Vec<T, A> {
    fn drop(&mut self) {
        let ptr_untagged = self.ptr_untagged();

        if let Some(ptr) = NonNull::new(ptr_untagged) {
            for i in 0..self.length() {
                unsafe {
                    ptr_untagged.add(i).drop_in_place();
                }
            }

            let ptr = ptr.cast();
            let capacity = self.capacity();
            let layout = Layout::array::<T>(capacity).unwrap();

            unsafe {
                self.allocator.deallocate(ptr, layout);
            }
        }
    }
}

#[cfg(test)]
mod test {
    use super::Vec;

    #[test]
    fn has_the_correct_size() {
        assert_eq!(std::mem::size_of::<Vec<u8>>(), 8);
    }

    #[test]
    fn works() {
        let mut vec = Vec::new();

        vec.push(1);
        vec.push(2);
        vec.push(3);

        assert_eq!(vec.capacity(), 4);
        assert_eq!(vec.length(), 3);

        vec[2] = 5;

        assert_eq!(vec.get(0), Some(&1));
        assert_eq!(vec.get(1), Some(&2));
        assert_eq!(vec.get(2), Some(&5));
        assert_eq!(vec.get(3), None);

        assert_eq!(vec[1], 2);
    }

    #[test]
    fn drops_elements_on_drop() {
        struct Element<'a>(&'a mut bool);

        impl<'a> Drop for Element<'a> {
            fn drop(&mut self) {
                *self.0 = true;
            }
        }

        let mut vec = Vec::new();
        let mut was_dropped = false;
        vec.push(Element(&mut was_dropped));
        drop(vec);

        assert_eq!(was_dropped, true);
    }

    #[test]
    fn drops_elements_exactly_once() {
        struct Element<'a>(&'a mut u64);

        impl<'a> Drop for Element<'a> {
            fn drop(&mut self) {
                *self.0 += 1;
            }
        }

        let mut vec = Vec::new();
        let mut a = 0;
        let mut b = 0;
        vec.push(Element(&mut a));
        vec.push(Element(&mut b));

        for elem in vec.into_iter().take(1) {
            // not part of the test
            assert_eq!(*elem.0, 0);
        }

        assert_eq!(a, 1);
        assert_eq!(b, 1);
    }

    #[test]
    #[should_panic]
    fn panics_on_oob() {
        let mut vec = Vec::new();

        vec.push(1);

        dbg!(vec[2]);
    }
}