sketchy-vec.rs

use alloc::alloc;

use core::{
    alloc::Layout,
    convert::{TryFrom, TryInto},
    fmt,
    mem::{self, ManuallyDrop},
    num::NonZeroU32,
    ops::{self, Drop},
    ptr,
    ptr::NonNull,
};

#[repr(C)]
struct NonEmptyVecInner<T> {
    len: u32,
    cap: NonZeroU32,
    rest: [T; 0],
}

#[repr(C)]
struct EmptyVecInner {
    len: u32,
    cap: u32,
}

static EMPTY: EmptyVecInner = EmptyVecInner { len: 0, cap: 0 };

union VecInner<T> {
    nonempty: NonNull<NonEmptyVecInner<T>>,
    empty: NonNull<EmptyVecInner>,
}

pub struct Vec<T> {
    ptr: VecInner<T>,
}

unsafe impl<T> Send for Vec<T> where T: Send {}
unsafe impl<T> Sync for Vec<T> where T: Sync {}

impl<T> Drop for Vec<T> {
    fn drop(&mut self) {
        self.drain();

        if self.capacity() != 0 {
            let old_layout = Layout::new::<u32>()
                .extend(Layout::new::<NonZeroU32>())
                .unwrap()
                .0
                .extend(Layout::array::<T>(self.capacity()).unwrap())
                .unwrap()
                .0;

            unsafe {
                alloc::dealloc(self.ptr.nonempty.cast().as_ptr(), old_layout);
            }
        }
    }
}

impl<T> Default for Vec<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> Clone for Vec<T>
where
    T: Clone,
{
    fn clone(&self) -> Self {
        let mut out = Vec::new();
        out.extend(self.iter().cloned());
        out
    }
}

impl<T> fmt::Debug for Vec<T>
where
    T: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let slice: &[T] = self.as_ref();
        slice.fmt(f)
    }
}

#[cfg(debug_assertions)]
macro_rules! assert_unchecked {
    ($cond:expr) => {
        debug_assert!($cond);
    };
}

#[cfg(not(debug_assertions))]
macro_rules! assert_unchecked {
    ($cond:expr) => {
        if !$cond {
            core::hint::unreachable_unchecked();
        }
    };
}

pub struct IntoIter<T> {
    vec: ManuallyDrop<Vec<T>>,
    iter: ManuallyDrop<DrainInner<T>>,
}

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

    fn next(&mut self) -> Option<T> {
        self.iter.next()
    }
}

impl<T> Drop for IntoIter<T> {
    fn drop(&mut self) {
        unsafe {
            ManuallyDrop::drop(&mut self.iter);
            ManuallyDrop::drop(&mut self.vec);
        }
    }
}

struct DrainInner<T> {
    ptr: NonNull<T>,
    end: NonNull<T>,
}

pub struct Drain<'a, T> {
    _marker: core::marker::PhantomData<&'a mut Vec<T>>,
    inner: DrainInner<T>,
}

impl<T> Iterator for Drain<'_, T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        self.inner.next()
    }
}

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

    fn next(&mut self) -> Option<T> {
        if self.ptr == self.end {
            None
        } else {
            unsafe {
                let ptr = self.ptr;
                self.ptr = NonNull::new_unchecked(self.ptr.as_ptr().offset(1));
                Some(ptr::read(ptr.as_ptr()))
            }
        }
    }
}

impl<T> Drop for DrainInner<T> {
    fn drop(&mut self) {
        if mem::needs_drop::<T>() {
            while let Some(_) = self.next() {}
        }
    }
}

impl<T> ops::Index<usize> for Vec<T> {
    type Output = T;

    fn index(&self, i: usize) -> &T {
        self.get(i).unwrap()
    }
}

impl<T> ops::IndexMut<usize> for Vec<T> {
    fn index_mut(&mut self, i: usize) -> &mut T {
        self.get_mut(i).unwrap()
    }
}

impl<'a, T> IntoIterator for Vec<T> {
    type Item = T;
    type IntoIter = IntoIter<T>;

    fn into_iter(mut self) -> Self::IntoIter {
        let slice: &mut [T] = self.as_mut();
        let ptr = slice.as_mut_ptr();
        let end = unsafe { ptr.offset(slice.len().try_into().unwrap()) };

        let (ptr, end) = unsafe { (NonNull::new_unchecked(ptr), NonNull::new_unchecked(end)) };

        IntoIter {
            vec: ManuallyDrop::new(self),
            iter: ManuallyDrop::new(DrainInner { ptr, end }),
        }
    }
}

impl<'a, T> IntoIterator for &'a Vec<T> {
    type Item = &'a T;
    type IntoIter = core::slice::Iter<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.as_ref().into_iter()
    }
}

impl<'a, T> IntoIterator for &'a mut Vec<T> {
    type Item = &'a mut T;
    type IntoIter = core::slice::IterMut<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.as_mut().into_iter()
    }
}

impl<T> AsRef<[T]> for Vec<T> {
    fn as_ref(&self) -> &[T] {
        if self.len() == 0 {
            &[]
        } else {
            unsafe {
                core::slice::from_raw_parts(
                    &*NonNull::<[T; 0]>::from(&self.ptr.nonempty.as_ref().rest)
                        .cast::<T>()
                        .as_ptr(),
                    self.len(),
                )
            }
        }
    }
}

impl<T> AsMut<[T]> for Vec<T> {
    fn as_mut(&mut self) -> &mut [T] {
        if self.len() == 0 {
            &mut []
        } else {
            unsafe {
                core::slice::from_raw_parts_mut(
                    &mut *NonNull::<[T; 0]>::from(&self.ptr.nonempty.as_ref().rest)
                        .cast::<T>()
                        .as_ptr(),
                    self.len(),
                )
            }
        }
    }
}

impl<T> Vec<T> {
    pub fn new() -> Self {
        Vec {
            ptr: VecInner {
                empty: NonNull::from(&EMPTY),
            },
        }
    }

    pub fn last(&self) -> Option<&T> {
        Some(unsafe { self.get_unchecked(self.len().checked_sub(1)?) })
    }

    pub fn last_mut(&mut self) -> Option<&mut T> {
        Some(unsafe { self.get_unchecked_mut(self.len().checked_sub(1)?) })
    }

    pub fn with_capacity(cap: usize) -> Self {
        let mut out = Self::new();
        out.reserve(cap);
        out
    }

    pub fn extend<I>(&mut self, iter: I)
    where
        I: IntoIterator<Item = T>,
    {
        let iter = iter.into_iter();
        self.reserve(iter.size_hint().0);
        for i in iter {
            self.push(i);
        }
    }

    pub fn clear(&mut self) {
        self.drain();
    }

    pub fn drain(&mut self) -> Drain<'_, T> {
        let slice: &mut [T] = self.as_mut();
        let ptr = slice.as_mut_ptr();
        let len = slice.len();
        let end = unsafe { ptr.offset(len.try_into().unwrap()) };

        let (ptr, end) = unsafe { (NonNull::new_unchecked(ptr), NonNull::new_unchecked(end)) };

        if len != 0 {
            unsafe {
                self.set_len(0);
            }
        }

        Drain {
            inner: DrainInner { ptr, end },
            _marker: Default::default(),
        }
    }

    pub unsafe fn set_len(&mut self, new_len: usize) {
        self.ptr.nonempty.as_mut().len = new_len.try_into().unwrap();
    }

    pub fn iter(&self) -> <&Self as IntoIterator>::IntoIter {
        self.into_iter()
    }

    pub fn iter_mut(&mut self) -> <&mut Self as IntoIterator>::IntoIter {
        self.into_iter()
    }

    pub fn get(&self, i: usize) -> Option<&T> {
        // This could be `static_assert`
        assert!(mem::align_of::<EmptyVecInner>() <= mem::align_of::<NonEmptyVecInner<T>>());

        self.as_ref().get(i)
    }

    pub fn get_mut(&mut self, i: usize) -> Option<&mut T> {
        // This could be `static_assert`
        assert!(mem::align_of::<EmptyVecInner>() <= mem::align_of::<NonEmptyVecInner<T>>());

        self.as_mut().get_mut(i)
    }

    pub unsafe fn get_unchecked(&self, i: usize) -> &T {
        // This could be `static_assert`
        assert!(mem::align_of::<EmptyVecInner>() <= mem::align_of::<NonEmptyVecInner<T>>());

        self.as_ref().get_unchecked(i)
    }

    pub unsafe fn get_unchecked_mut(&mut self, i: usize) -> &mut T {
        // This could be `static_assert`
        assert!(mem::align_of::<EmptyVecInner>() <= mem::align_of::<NonEmptyVecInner<T>>());

        self.as_mut().get_unchecked_mut(i)
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn len(&self) -> usize {
        assert!(mem::align_of::<EmptyVecInner>() <= mem::align_of::<NonEmptyVecInner<T>>());

        unsafe { self.ptr.empty.as_ref().len as usize }
    }

    pub fn capacity(&self) -> usize {
        assert!(mem::align_of::<EmptyVecInner>() <= mem::align_of::<NonEmptyVecInner<T>>());

        unsafe { self.ptr.empty.as_ref().cap as usize }
    }

    pub fn push(&mut self, val: T) {
        assert!(mem::align_of::<EmptyVecInner>() <= mem::align_of::<NonEmptyVecInner<T>>());

        let len = self.len();
        let cap = self.capacity();

        unsafe {
            assert_unchecked!(len <= cap);
        }

        if len == cap {
            let new_len = (cap << 1).max(4);
            assert!(new_len > len);
            self.reserve((new_len - len) as usize);
        }

        unsafe {
            self.push_unchecked(val);
        }
    }

    unsafe fn push_unchecked(&mut self, val: T) {
        let len = self.len();

        debug_assert!(len < self.capacity(), "{:?}", (len, self.capacity()));

        ptr::write(
            NonNull::<[T; 0]>::from(&self.ptr.nonempty.as_ref().rest)
                .cast::<T>()
                .as_ptr()
                .offset(len.try_into().unwrap()),
            val,
        );

        self.set_len(len + 1);
    }

    pub fn reserve(&mut self, count: usize) {
        assert!(mem::align_of::<EmptyVecInner>() <= mem::align_of::<NonEmptyVecInner<T>>());

        if count == 0 {
            return;
        }

        let len = self.len();
        let cap = self.capacity();
        let new_cap = u32::try_from(len as usize + count).unwrap();

        unsafe {
            assert_unchecked!(len <= cap);
        }

        if cap >= new_cap as usize {
            return;
        }

        let layout = Layout::new::<u32>()
            .extend(Layout::new::<NonZeroU32>())
            .unwrap()
            .0
            .extend(Layout::array::<T>(new_cap as usize).unwrap())
            .unwrap()
            .0;

        let mut this = ManuallyDrop::new(mem::replace(self, Vec::new()));

        let new_ptr = unsafe {
            NonNull::new(if this.ptr.empty == NonNull::from(&EMPTY) {
                alloc::alloc(layout)
            } else {
                let ptr = this.ptr.nonempty.as_ptr() as *mut u8;

                let old_layout = Layout::new::<u32>()
                    .extend(Layout::new::<NonZeroU32>())
                    .unwrap()
                    .0
                    .extend(Layout::array::<T>(cap as usize).unwrap())
                    .unwrap()
                    .0;

                let new_ptr = alloc::realloc(ptr, old_layout, layout.size());

                new_ptr
            })
            .unwrap()
            .cast::<NonEmptyVecInner<T>>()
        };

        this.ptr = VecInner { nonempty: new_ptr };
        unsafe {
            this.ptr.nonempty.as_mut().cap = NonZeroU32::new(new_cap).unwrap();
            this.ptr.nonempty.as_mut().len = len as u32;
        }

        *self = ManuallyDrop::into_inner(this);
    }
}

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

    #[test]
    fn it_works() {
        let mut v = Vec::new();
        v.push(1u32);
        v.push(2u32);
        v.push(3u32);
        v.push(4u32);
        v.push(5u32);
        v.push(6u32);
        v.push(7u32);

        assert_eq!(v.capacity(), 8);
        assert_eq!(v.len(), 7);
        for i in 0..v.len() {
            assert_eq!(v.get(i), Some(i as u32 + 1).as_ref());
        }
    }
}