diff --git a/src/lib.rs b/src/lib.rs index c2c41a9..81f67ae 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,63 +1,153 @@ +#![warn(clippy::pedantic)] + #[cfg(test)] mod tests; use std::alloc::Layout; +use std::hint::unreachable_unchecked; use std::mem::ManuallyDrop; use std::ptr::NonNull; +/// Stores the capacity of a `Vec` for later reuse as a `Vec` where `V` shares the same `Layout` as `U`. pub struct TypeErasedVec { + /// Pointer to the underlying allocation ptr: NonNull, - len: usize, - cap: usize, + /// Size of the underlying allocation + capacity: usize, + /// The Layout the capacity was allocated with. We need this to confirm that + /// any future conversion back to a `Vec` use the correct `Layout`. layout: Layout, - // Store a function pointer to the original type's drop logic. - // This ensures elements and memory are cleaned up if this struct is dropped. - drop_impl: unsafe fn(NonNull, usize, usize), + /// Function pointer to the original type's drop logic. + /// While we never drop any elements with it, we need to drop the capacity + /// by converting it back to an (empty) Vec of a compatibly type. + drop_capacity: unsafe fn(NonNull, usize), } +// SAFETY: TypeErasedVec only holds an empty memory allocation and a stateless function pointer. +// All elements are explicitly dropped during construction so no instances of the original type exist. +// Transferring ownership of this uninitialized capacity across threads is safe because the underlying +// global allocator is thread-safe and there is no data to cause data races. +unsafe impl Send for TypeErasedVec {} + +// SAFETY: TypeErasedVec contains no interior mutability. +// It only exposes immutable metadata regarding the underlying allocation when accessed via a shared reference. +// Sharing a reference to this empty allocation across threads cannot cause data races or undefined behavior. +unsafe impl Sync for TypeErasedVec {} + impl TypeErasedVec { - pub fn new(vec: Vec) -> Self { + /// Type erase the underlying Capacity of a `Vec` remembering the `Layout` it was allocated with. + /// Any remaining elements in the `Vec` will be dropped. + /// Conversion back to a `Vec` is only allowed for types with the same Layout. + #[must_use] + pub fn new(mut vec: Vec) -> Self { + // Define a cleanup function bound to the original type. + unsafe fn drop_vec_capacity(ptr: NonNull, cap: usize) { + let _ = unsafe { Vec::from_raw_parts(ptr.as_ptr().cast::(), 0, cap) }; + } + let layout = Layout::new::(); + // Ensure the vector is empty. We don't do type casts, we just reuse capacity. + vec.clear(); + // Deconstruct the original vector into its raw components. - let (ptr, len, cap) = vec.into_raw_parts(); + let (ptr, len, capacity) = vec.into_raw_parts(); + + debug_assert_eq!(len, 0); // Vec guarantees its pointer is never null, even when capacity is zero. - let ptr = unsafe { NonNull::new_unchecked(ptr as *mut u8) }; - - // Define a cleanup function bound to the original type. - unsafe fn drop_vec(ptr: NonNull, len: usize, cap: usize) { - let _ = Vec::from_raw_parts(ptr.as_ptr() as *mut T, len, cap); - } + let ptr = unsafe { NonNull::new_unchecked(ptr.cast::()) }; Self { ptr, - len, - cap, + capacity, layout, - drop_impl: drop_vec::, + drop_capacity: drop_vec_capacity::, } } + /// Get the layout of the underlying capacity + #[must_use] + pub fn layout(&self) -> Layout { + self.layout + } + + /// Get the underlying capacity in units of the `Layout` size + #[must_use] + pub fn capacity(&self) -> usize { + self.capacity + } + + /// Get the underlying capacity in bytes + #[must_use] + pub fn capacity_bytes(&self) -> usize { + self.capacity * self.layout.size() + } + + /// Convert the capacity of the erased `Vec` into a `Vec`. + /// + /// # Errors + /// if `T` does not have the same `Layout` as the underlying capacity + /// in this case the erased capacity will be returned in the `Err`. + pub fn try_into_vec(self) -> Result, Self> { + if self.layout == Layout::new::() { + // Wrap self in ManuallyDrop to bypass our custom Drop implementation. + // This transfers ownership of the memory to the new Vec. + let this = ManuallyDrop::new(self); + let vec = unsafe { Vec::from_raw_parts(this.ptr.as_ptr().cast(), 0, this.capacity) }; + + Ok(vec) + } else { + Err(self) + } + } + + /// Convert the capacity of the erased `Vec` into a `Vec`. + /// + /// # Panics + /// If `T` does not have the same `Layout` as the underlying capacity. + #[must_use] pub fn into_vec(self) -> Vec { - assert_eq!( - Layout::new::(), - self.layout, - "Target type layout must exactly match the erased layout" - ); + match self.try_into_vec::() { + Ok(vec) => vec, + Err(this) => { + panic!( + "Target type layout must exactly match the erased layout. Capacity is reserved for {:?} but {} has {:?}", + this.layout, + std::any::type_name::(), + Layout::new::() + ) + } + } + } - // Wrap self in ManuallyDrop to bypass our custom Drop implementation. - // This transfers ownership of the memory to the new Vec. - let md = ManuallyDrop::new(self); - - unsafe { Vec::from_raw_parts(md.ptr.as_ptr() as *mut T, md.len, md.cap) } + /// Convert the capacity of the erased `Vec` into a `Vec`. + /// + /// # Safety + /// `T` must have the same `Layout` as the underlying capacity. + #[must_use] + pub unsafe fn into_vec_unchecked(self) -> Vec { + match self.try_into_vec::() { + Ok(vec) => vec, + Err(this) => { + if cfg!(debug_assertions) { + unreachable!( + "Calling `into_vec_unchecked` with an incompatible layout is UB! Target type layout must exactly match the erased layout. Capacity is reserved for {:?} but {} has {:?}", + this.layout, + std::any::type_name::(), + Layout::new::() + ) + } + unsafe { unreachable_unchecked() } + } + } } } impl Drop for TypeErasedVec { fn drop(&mut self) { unsafe { - (self.drop_impl)(self.ptr, self.len, self.cap); + (self.drop_capacity)(self.ptr, self.capacity); } } } diff --git a/src/tests.rs b/src/tests.rs index 48c3edb..691e8f2 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -1,135 +1,165 @@ -#[cfg(test)] -mod tests { - use crate::TypeErasedVec; +use super::*; +use std::rc::Rc; - use super::*; - use std::cell::Cell; - use std::rc::Rc; +#[test] +fn test_basic_roundtrip() { + let mut v = Vec::::with_capacity(10); + v.push(1); + v.push(2); + let cap = v.capacity(); - #[repr(C)] - #[derive(Debug, PartialEq)] - struct TypeA { - a: u32, - b: u16, - c: u16, - } + let erased = TypeErasedVec::new(v); + let recovered: Vec = erased.into_vec(); - #[repr(C)] - #[derive(Debug, PartialEq)] - struct TypeB { - x: u64, - } + assert_eq!(recovered.len(), 0); + assert_eq!(recovered.capacity(), cap); +} + +#[test] +fn test_compatible_types() { + let mut v = Vec::::with_capacity(42); + v.push(100); + let cap = v.capacity(); + + let erased = TypeErasedVec::new(v); + + // Convert to a different type with the exact same Layout + let recovered: Vec = erased.try_into_vec().ok().unwrap(); + + assert_eq!(recovered.len(), 0); + assert_eq!(recovered.capacity(), cap); +} + +#[test] +fn test_incompatible_size() { + let v = Vec::::with_capacity(10); + let erased = TypeErasedVec::new(v); + + let result: Result, TypeErasedVec> = erased.try_into_vec(); + + // Result should be Err containing the original erased vec + let Err(erased) = result else { + panic!("Expected conversion from u64 to u8 to fail"); + }; + + assert_eq!(erased.capacity(), 10); + assert_eq!(erased.into_vec::().capacity(), 10); +} + +#[test] +fn test_incompatible_alignment() { + #[repr(align(16))] + struct Align16(#[allow(unused)] u8); + + #[repr(align(8))] + struct Align8(#[allow(unused)] u8); - // A helper to track when values are dropped - // This is crucial for verifying that TypeErasedVec does not leak memory or skip destructors - #[derive(Debug)] - struct DropTracker { - counter: Rc>, - } + let v = Vec::::with_capacity(10); + let erased = TypeErasedVec::new(v); - impl Drop for DropTracker { - fn drop(&mut self) { - self.counter.set(self.counter.get() + 1); - } - } + // Sizes might be compatible or both be wrapped in padding, but alignment differs + let result: Result, TypeErasedVec> = erased.try_into_vec(); - #[test] - fn test_successful_conversion() { - let original = vec![TypeA { a: 1, b: 2, c: 3 }, TypeA { a: 4, b: 5, c: 6 }]; + let Err(erased) = result else { + panic!("Expected conversion from u64 to u8 to fail"); + }; - let erased = TypeErasedVec::new(original); - let converted = erased.into_vec::(); + assert_eq!(erased.capacity(), 10); + assert_eq!(erased.into_vec::().capacity(), 10); +} - assert_eq!(converted.len(), 2); +#[test] +#[should_panic(expected = "Target type layout must exactly match")] +fn test_into_vec_panic_on_mismatch() { + let v = Vec::::with_capacity(10); + let erased = TypeErasedVec::new(v); - // The exact bit pattern of TypeA {1, 2, 3} depends on endianness - // Miri will ensure this read is memory-safe regardless of the values inside - } + // This should panic due to Layout mismatch + let _panic: Vec = erased.into_vec(); +} - #[test] - fn test_zero_capacity_vec() { - // Vec::new() does not allocate - // The pointer is dangling but non-null - let original: Vec = Vec::new(); +#[test] +fn test_elements_are_dropped() { + let counter = Rc::new(()); + let mut v = Vec::new(); - let erased = TypeErasedVec::new(original); - let converted = erased.into_vec::(); + v.push(Rc::clone(&counter)); + v.push(Rc::clone(&counter)); - assert_eq!(converted.capacity(), 0); - assert_eq!(converted.len(), 0); - } + assert_eq!(Rc::strong_count(&counter), 3); - #[test] - fn test_zero_sized_types() { - // ZSTs do not allocate memory but the length must be tracked correctly - let original = vec![(), (), ()]; + let erased = TypeErasedVec::new(v); - let erased = TypeErasedVec::new(original); - let converted = erased.into_vec::<()>(); + // Elements should have been dropped during `TypeErasedVec::new` by `vec.clear()` + assert_eq!(Rc::strong_count(&counter), 1); - assert_eq!(converted.len(), 3); - } + std::mem::drop(erased); - #[test] - fn test_erased_vec_drop_cleans_up_elements() { - let drop_count = Rc::new(Cell::new(0)); + assert_eq!(Rc::strong_count(&counter), 1); +} - let original = vec![ - DropTracker { - counter: Rc::clone(&drop_count), - }, - DropTracker { - counter: Rc::clone(&drop_count), - }, - DropTracker { - counter: Rc::clone(&drop_count), - }, - ]; +#[test] +fn test_zst_handling() { + let mut v = Vec::<()>::with_capacity(10); + v.push(()); + v.push(()); - let erased = TypeErasedVec::new(original); + let erased = TypeErasedVec::new(v); + let recovered: Vec<()> = erased.into_vec(); - // Dropping the erased container must trigger the original type's drop logic - drop(erased); + assert_eq!(recovered.len(), 0); +} - assert_eq!(drop_count.get(), 3); - } +#[test] +fn test_zero_capacity() { + let v = Vec::::new(); + let erased = TypeErasedVec::new(v); - #[test] - fn test_capacity_is_preserved() { - let mut original = Vec::with_capacity(42); - original.push(TypeA { a: 0, b: 0, c: 0 }); + let recovered: Vec = erased.into_vec(); - let erased = TypeErasedVec::new(original); - let converted = erased.into_vec::(); + assert_eq!(recovered.len(), 0); + assert_eq!(recovered.capacity(), 0); +} - assert_eq!(converted.capacity(), 42); - assert_eq!(converted.len(), 1); - } +#[test] +fn test_drop_erased_memory_leak() { + let v = Vec::::with_capacity(100); + let erased = TypeErasedVec::new(v); - #[test] - #[should_panic(expected = "Target type layout must exactly match the erased layout")] - fn test_panic_on_size_mismatch() { - let original = vec![1u32, 2u32]; - let erased = TypeErasedVec::new(original); + // Miri will flag this test as failing if `erased` doesn't properly deallocate + // the underlying memory upon being dropped. + drop(erased); +} - // u64 has a different size than u32 - // This must panic to prevent memory corruption and Miri errors - let _converted = erased.into_vec::(); - } +#[test] +fn test_into_vec_unchecked_success() { + let mut v = Vec::with_capacity(15); + v.push(42); + let cap = v.capacity(); - #[test] - #[should_panic(expected = "Target type layout must exactly match the erased layout")] - fn test_panic_on_alignment_mismatch() { - #[repr(align(8))] - struct Aligned8([u8; 8]); + let erased = TypeErasedVec::new(v); - #[repr(align(1))] - struct Aligned1([u8; 8]); + let recovered: Vec = unsafe { erased.into_vec_unchecked() }; - let original = vec![Aligned8([0; 8])]; - let erased = TypeErasedVec::new(original); + assert_eq!(recovered.len(), 0); + assert_eq!(recovered.capacity(), cap); +} - // Both types are 8 bytes but they have different alignments - let _converted = erased.into_vec::(); +#[test] +fn test_complex_struct_drop() { + struct Droppy { + _a: String, + _b: Vec, } + + let mut v = Vec::with_capacity(5); + v.push(Droppy { + _a: String::from("Hello"), + _b: vec![1, 2, 3], + }); + + let erased = TypeErasedVec::new(v); + + // Ensures `TypeErasedVec` properly drops via the generic drop function + drop(erased); }