diff --git a/src/tinyvec.rs b/src/tinyvec.rs index 45cc8e2..7f07253 100644 --- a/src/tinyvec.rs +++ b/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 { 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") } @@ -63,32 +59,35 @@ impl TinyVec { 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::(), 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::() <= 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 - .add(old_len) - .copy_from(slice.as_ptr(), slice.len()); - }, - } + false => unsafe { self.inner.heap.data }, + }; - unsafe { self.inner.len.len += slice.len() as u8 } + unsafe { + ptr.cast::() + .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) { @@ -96,20 +95,24 @@ impl TinyVec { 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::(); + 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::().write(v) }, - } + false => unsafe { self.inner.heap.data }, + }; - unsafe { self.inner.len.len += 1 }; + unsafe { + ptr.cast::().add(self.len()).write(v); + self.inner.len.len += 1 + }; } fn spill(&mut self) { @@ -117,7 +120,7 @@ impl TinyVec { 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 Drop for TinyVec { 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);