diff --git a/src/lib.rs b/src/lib.rs index a977fa4..2fb5dc1 100644 --- a/src/lib.rs +++ b/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)] diff --git a/src/tinyvec.rs b/src/tinyvec.rs new file mode 100644 index 0000000..45cc8e2 --- /dev/null +++ b/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, + stack: TinyStackVec, +} + +pub struct TinyVec { + inner: TinyVecInner, + _ph: PhantomData, +} + +impl TinyVec { + pub fn new() -> Self { + const { + if std::mem::size_of::() != 1 { + panic!("sizeof T must be 1") + } + if std::mem::align_of::() != 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::().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 std::ops::Deref for TinyVec { + 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 std::ops::DerefMut for TinyVec { + 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 Extend for TinyVec { + fn extend>(&mut self, iter: I) { + for v in iter { + self.push(v); + } + } +} + +impl FromIterator for TinyVec { + fn from_iter>(iter: I) -> Self { + let mut s = Self::new(); + s.extend(iter); + s + } +} + +pub struct IntoIter { + index: usize, + v: TinyVec, +} + +impl Iterator for IntoIter { + type Item = T; + + fn next(&mut self) -> Option { + match self.v.get(self.index) { + None => None, + Some(&e) => { + self.index += 1; + Some(e) + } + } + } +} + +impl IntoIterator for TinyVec { + type Item = T; + + type IntoIter = IntoIter; + + fn into_iter(self) -> Self::IntoIter { + IntoIter { index: 0, v: self } + } +} + +impl<'a, T: Copy> IntoIterator for &'a TinyVec { + type Item = T; + + type IntoIter = Copied>; + + fn into_iter(self) -> Self::IntoIter { + self.iter().copied() + } +} + +impl Drop for TinyVec { + fn drop(&mut self) { + if !self.is_stack() { + unsafe { + std::alloc::dealloc(self.inner.heap.data, HEAP_LAYOUT); + } + } + } +} + +impl Default for TinyVec { + 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::::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::::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::::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::::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()); + } +}