Add a tinyvec

src/lib.rs

@@ -3,6 +3,8 @@ use std::path::PathBuf;
 use bstr::{BString, ByteSlice};
 use clap::Parser;
 
+pub mod tinyvec;
+
 #[derive(Parser)]
 struct Args {
     #[arg(short, long)]

src/tinyvec.rs

@@ -0,0 +1,327 @@
+use std::{alloc::Layout, iter::Copied, marker::PhantomData, mem::ManuallyDrop};
+
+const STACK_LEN: usize = std::mem::size_of::<*mut u8>() * 2 - 1;
+const HEAP_LAYOUT: Layout = unsafe { Layout::from_size_align_unchecked(127, 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::size_of::<T>() != 1 {
+                panic!("sizeof T must be 1")
+            }
+            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 }
+    }
+
+    pub fn extend_from_slice(&mut self, slice: &[u8]) {
+        let old_len = self.len();
+        let new_len = old_len + slice.len();
+        if new_len > 127 {
+            panic!("Out of capacity")
+        }
+
+        match self.is_stack() {
+            true if new_len <= STACK_LEN => unsafe {
+                self.inner.stack.data[old_len..new_len].copy_from_slice(slice);
+            },
+            true => {
+                self.spill();
+                self.extend_from_slice(slice);
+                return;
+            }
+            false => unsafe {
+                self.inner
+                    .heap
+                    .data
+                    .add(old_len)
+                    .copy_from(slice.as_ptr(), slice.len());
+            },
+        }
+
+        unsafe { self.inner.len.len += slice.len() as u8 }
+    }
+
+    pub fn push(&mut self, v: T) {
+        if self.len() == 127 {
+            panic!("Out of capacity");
+        }
+
+        match self.is_stack() {
+            true => {
+                if self.len() == STACK_LEN {
+                    self.spill();
+                    self.push(v);
+                    return;
+                } else {
+                    unsafe { *(&raw mut self.inner.stack.data[self.len()] as *mut T) = v }
+                }
+            }
+            false => unsafe { self.inner.heap.data.add(self.len()).cast::<T>().write(v) },
+        }
+
+        unsafe { 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(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::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, 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]);
+    }
+
+    #[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);
+    }
+
+    #[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);
+    }
+
+    #[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 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..50).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..50).collect_vec();
+        let mut a = TinyVec::new();
+        a.extend_from_slice(&values);
+
+        assert_eq!(values, (&a).into_iter().collect_vec());
+    }
+}