Allow tinyvecs to be of any type with alignement 1

src/tinyvec.rs

@@ -1,7 +1,6 @@
 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 {
@@ -36,9 +35,6 @@ pub struct TinyVec<T: Copy> {
 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")
             }
@@ -63,32 +59,35 @@ impl<T: Copy> TinyVec<T> {
         unsafe { (self.inner.len.len & 127) as usize }
     }
 
-    pub fn extend_from_slice(&mut self, slice: &[u8]) {
+    fn heap_layout() -> Layout {
+        Layout::from_size_align(127 * size_of::<T>(), 1).unwrap()
+    }
+
+    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")
         }
 
-        match self.is_stack() {
-            true if new_len <= STACK_LEN => unsafe {
-                self.inner.stack.data[old_len..new_len].copy_from_slice(slice);
+        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
+            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
         }
-
-        unsafe { self.inner.len.len += slice.len() as u8 }
     }
 
     pub fn push(&mut self, v: T) {
@@ -96,20 +95,24 @@ impl<T: Copy> TinyVec<T> {
             panic!("Out of capacity");
         }
 
-        match self.is_stack() {
+        let ptr = match self.is_stack() {
             true => {
-                if self.len() == STACK_LEN {
+                let new_stack_len = (self.len() + 1) * size_of::<T>();
+                if new_stack_len > STACK_LEN {
                     self.spill();
                     self.push(v);
                     return;
                 } else {
-                    unsafe { *(&raw mut self.inner.stack.data[self.len()] as *mut T) = v }
+                    unsafe { self.inner.stack.data.as_mut_ptr() }
                 }
             }
-            false => unsafe { self.inner.heap.data.add(self.len()).cast::<T>().write(v) },
-        }
+            false => unsafe { self.inner.heap.data },
+        };
 
-        unsafe { self.inner.len.len += 1 };
+        unsafe {
+            ptr.cast::<T>().add(self.len()).write(v);
+            self.inner.len.len += 1
+        };
     }
 
     fn spill(&mut self) {
@@ -117,7 +120,7 @@ impl<T: Copy> TinyVec<T> {
         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 = std::alloc::alloc(Self::heap_layout());
             (*self.inner.heap)
                 .data
                 .copy_from(&data as *const _, size_of_val(&data));
@@ -221,7 +224,7 @@ 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);
+                std::alloc::dealloc(self.inner.heap.data, Self::heap_layout());
             }
         }
     }
@@ -256,6 +259,12 @@ mod test {
         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]
@@ -269,6 +278,16 @@ mod test {
         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]
@@ -282,6 +301,21 @@ mod test {
         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]
@@ -309,7 +343,7 @@ mod test {
 
     #[test]
     pub fn into_iter() {
-        let values = (0..50).collect_vec();
+        let values = (0..50u8).collect_vec();
         let mut a = TinyVec::new();
         a.extend_from_slice(&values);
 
@@ -318,7 +352,7 @@ mod test {
 
     #[test]
     pub fn into_ref_iter() {
-        let values = (0..50).collect_vec();
+        let values = (0..50u8).collect_vec();
         let mut a = TinyVec::new();
         a.extend_from_slice(&values);