aoc_2025/src/tinyvec.rs

use std::{alloc::Layout, iter::Copied, marker::PhantomData, mem::ManuallyDrop};

const STACK_LEN: usize = std::mem::size_of::<*mut u8>() * 2 - 1;

#[repr(C)]
struct TinyHeapVec {
    len: u8,
    data: *mut u8,
}

#[derive(Clone, Copy)]
#[repr(C)]
struct TinyStackVec {
    len: u8,
    data: [u8; STACK_LEN],
}

#[derive(Clone, Copy)]
#[repr(C)]
struct TinyVecLen {
    len: u8,
}

union TinyVecInner {
    len: TinyVecLen,
    heap: ManuallyDrop<TinyHeapVec>,
    stack: TinyStackVec,
}

pub struct TinyVec<T: Copy> {
    inner: TinyVecInner,
    _ph: PhantomData<T>,
}

impl<T: Copy> TinyVec<T> {
    pub fn new() -> Self {
        const {
            if std::mem::align_of::<T>() != 1 {
                panic!("alignof T must be 1")
            }
        }

        Self {
            inner: TinyVecInner {
                stack: TinyStackVec {
                    len: 0,
                    data: [0; STACK_LEN],
                },
            },
            _ph: PhantomData,
        }
    }

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

    pub fn len(&self) -> usize {
        unsafe { (self.inner.len.len & 127) as usize }
    }

    fn heap_layout() -> Layout {
        Layout::from_size_align(127 * size_of::<T>(), 1).unwrap()
    }

    pub fn insert_at(&mut self, at: usize, v: T) {
        if at > self.len() {
            panic!("Out of bound access")
        }

        if at == self.len() {
            return self.push(v);
        }

        let ptr = match self.is_stack() {
            true => {
                let new_stack_len = (self.len() + 1) * size_of::<T>();
                if new_stack_len > STACK_LEN {
                    self.spill();
                    self.insert_at(at, v);
                    return;
                } else {
                    unsafe { self.inner.stack.data.as_mut_ptr() }
                }
            }
            false => unsafe { self.inner.heap.data },
        };

        unsafe {
            let place = ptr.cast::<T>().add(at);
            place.add(1).copy_from(place, self.len() - at);
            place.write(v);
            self.inner.len.len += 1
        }
    }

    pub fn extend_from_slice(&mut self, slice: &[T]) {
        let old_len = self.len();
        let new_len = old_len + slice.len();
        if new_len > 127 {
            panic!("Out of capacity")
        }

        let ptr = match self.is_stack() {
            true if new_len * size_of::<T>() <= STACK_LEN => unsafe {
                self.inner.stack.data.as_mut_ptr()
            },
            true => {
                self.spill();
                self.extend_from_slice(slice);
                return;
            }
            false => unsafe { self.inner.heap.data },
        };

        unsafe {
            ptr.cast::<T>()
                .add(old_len)
                .copy_from(slice.as_ptr(), slice.len());
            self.inner.len.len += slice.len() as u8
        }
    }

    pub fn push(&mut self, v: T) {
        if self.len() == 127 {
            panic!("Out of capacity");
        }

        let ptr = match self.is_stack() {
            true => {
                let new_stack_len = (self.len() + 1) * size_of::<T>();
                if new_stack_len > STACK_LEN {
                    self.spill();
                    self.push(v);
                    return;
                } else {
                    unsafe { self.inner.stack.data.as_mut_ptr() }
                }
            }
            false => unsafe { self.inner.heap.data },
        };

        unsafe {
            ptr.cast::<T>().add(self.len()).write(v);
            self.inner.len.len += 1
        };
    }

    fn spill(&mut self) {
        assert!(self.is_stack());
        unsafe {
            let data = self.inner.stack.data;
            self.inner.len.len |= 128;
            (*self.inner.heap).data = std::alloc::alloc(Self::heap_layout());
            (*self.inner.heap)
                .data
                .copy_from(&data as *const _, size_of_val(&data));
        }
    }

    fn is_stack(&self) -> bool {
        let len = unsafe { self.inner.len.len } as usize;

        len & 128 == 0
    }
}

impl<T: Copy + std::fmt::Debug> std::fmt::Debug for TinyVec<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_list().entries(self).finish()
    }
}

impl<T: Copy> std::ops::Deref for TinyVec<T> {
    type Target = [T];

    fn deref(&self) -> &Self::Target {
        match self.is_stack() {
            true => unsafe {
                core::slice::from_raw_parts(self.inner.stack.data.as_ptr() as *const _, self.len())
            },
            false => unsafe {
                core::slice::from_raw_parts(self.inner.heap.data as *const _, self.len())
            },
        }
    }
}

impl<T: Copy> std::ops::DerefMut for TinyVec<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        match self.is_stack() {
            true => unsafe {
                core::slice::from_raw_parts_mut(
                    self.inner.stack.data.as_ptr() as *mut _,
                    self.len(),
                )
            },
            false => unsafe {
                core::slice::from_raw_parts_mut(self.inner.heap.data as *mut _, self.len())
            },
        }
    }
}

impl<T: Copy> Extend<T> for TinyVec<T> {
    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
        for v in iter {
            self.push(v);
        }
    }
}

impl<T: Copy> FromIterator<T> for TinyVec<T> {
    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
        let mut s = Self::new();
        s.extend(iter);
        s
    }
}

pub struct IntoIter<T: Copy> {
    index: usize,
    v: TinyVec<T>,
}

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

    fn next(&mut self) -> Option<Self::Item> {
        match self.v.get(self.index) {
            None => None,
            Some(&e) => {
                self.index += 1;
                Some(e)
            }
        }
    }
}

impl<T: Copy> IntoIterator for TinyVec<T> {
    type Item = T;

    type IntoIter = IntoIter<T>;

    fn into_iter(self) -> Self::IntoIter {
        IntoIter { index: 0, v: self }
    }
}

impl<'a, T: Copy> IntoIterator for &'a TinyVec<T> {
    type Item = T;

    type IntoIter = Copied<std::slice::Iter<'a, T>>;

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

impl<T: Copy> Drop for TinyVec<T> {
    fn drop(&mut self) {
        if !self.is_stack() {
            unsafe {
                std::alloc::dealloc(self.inner.heap.data, Self::heap_layout());
            }
        }
    }
}

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

#[cfg(test)]
mod test {
    use itertools::Itertools;

    use crate::tinyvec::STACK_LEN;

    use super::TinyVec;

    #[test]
    pub fn empty() {
        let a = TinyVec::<u8>::new();
        assert_eq!(a.len(), 0);
        assert!(a.is_empty());
        assert_eq!(&*a, &[]);
    }

    #[test]
    pub fn push_stack() {
        let mut a = TinyVec::new();
        a.push(42u8);
        assert_eq!(a.len(), 1);
        assert!(!a.is_empty());
        assert_eq!(&*a, &[42]);

        let mut a = TinyVec::new();
        a.push([42u8, 43]);
        assert_eq!(a.len(), 1);
        assert!(!a.is_empty());
        assert_eq!(&*a, &[[42, 43]]);
    }

    #[test]
    pub fn extend_stack() {
        let values = (0..STACK_LEN as u8).collect_vec();
        let mut a = TinyVec::<u8>::new();
        a.extend_from_slice(&values);
        assert_eq!(&*a, &values);

        let mut a = TinyVec::new();
        a.push(values[0]);
        a.extend_from_slice(&values[1..]);
        assert_eq!(&*a, &values);

        let values = [[43u8, 45], [46, 47]];
        let mut a = TinyVec::new();
        a.extend_from_slice(&values);
        assert_eq!(&*a, &values);

        let mut a = TinyVec::new();
        a.push(values[0]);
        a.extend_from_slice(&values[1..]);
        assert_eq!(&*a, &values);
    }

    #[test]
    pub fn extend_spill() {
        let values = (0..STACK_LEN as u8 + 1).collect_vec();
        let mut a = TinyVec::<u8>::new();
        a.extend_from_slice(&values);
        assert_eq!(&*a, &values);

        let mut a = TinyVec::new();
        a.push(values[0]);
        a.extend_from_slice(&values[1..]);
        assert_eq!(&*a, &values);

        let values = [
            [43u8, 45, 46, 47],
            [48, 49, 50, 51],
            [52, 53, 54, 55],
            [56, 57, 58, 59],
        ];
        let mut a = TinyVec::new();
        a.extend_from_slice(&values);
        assert_eq!(&*a, &values);

        let mut a = TinyVec::new();
        a.push(values[0]);
        a.extend_from_slice(&values[1..]);
        assert_eq!(&*a, &values);
    }

    #[test]
    pub fn extend_heap() {
        let values = (0..127).collect_vec();
        let mut a = TinyVec::<u8>::new();
        a.extend_from_slice(&values[0..10]);
        a.extend_from_slice(&values[10..]);
        assert_eq!(&*a, values);
    }

    #[test]
    pub fn insert_stack() {
        let mut a = TinyVec::<[u8; 2]>::new();
        a.push([1, 2]);
        a.push([3, 4]);
        a.insert_at(1, [5, 6]);
        assert_eq!(&*a, &[[1, 2], [5, 6], [3, 4]]);
    }

    #[test]
    pub fn insert_heap() {
        let mut values = vec![[0, 0]; 16];

        let mut a = TinyVec::<[u8; 2]>::new();
        a.extend_from_slice(&values);

        values.insert(1, [5, 6]);
        a.insert_at(1, [5, 6]);

        assert_eq!(&*a, &values);
    }

    #[test]
    pub fn insert_spill() {
        let mut values = vec![[0, 0]; 7];

        let mut a = TinyVec::<[u8; 2]>::new();
        a.extend_from_slice(&values);

        values.insert(1, [5, 6]);
        a.insert_at(1, [5, 6]);

        assert_eq!(&*a, &values);
    }

    #[test]
    pub fn push_spill() {
        let values = (0..STACK_LEN as u8 + 1).collect_vec();

        let mut a = TinyVec::new();
        for &i in &values {
            a.push(i);
        }

        assert_eq!(a.len(), values.len());
        assert!(!a.is_empty());
        assert_eq!(&*a, &values);
    }

    #[test]
    pub fn into_iter() {
        let values = (0..50u8).collect_vec();
        let mut a = TinyVec::new();
        a.extend_from_slice(&values);

        assert_eq!(values, a.into_iter().collect_vec());
    }

    #[test]
    pub fn into_ref_iter() {
        let values = (0..50u8).collect_vec();
        let mut a = TinyVec::new();
        a.extend_from_slice(&values);

        assert_eq!(values, (&a).into_iter().collect_vec());
    }
}