diff --git a/src/lib.rs b/src/lib.rs index 2810b7fd..5434d44f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -142,7 +142,8 @@ where R: core::ops::RangeBounds { #[repr(C)] pub struct SmallVec { len: TaggedLen, - raw: RawSmallVec + raw: RawSmallVec, + allocator: A } unsafe impl Send for SmallVec {} @@ -252,7 +253,8 @@ pub struct IntoIter { // `end` decides whether the data lives on the heap or not // // The members from begin..end are initialized - raw: RawSmallVec, + raw: RawSmallVec, + allocator: A, begin: usize, end: TaggedLen } @@ -371,7 +373,8 @@ impl SmallVec { // SAFETY: all the members in 0..S are initialized Self { len: TaggedLen::new(S, false), - raw: RawSmallVec::new_inline(buf, Global) + raw: RawSmallVec::new_inline(buf), + allocator: Global } } @@ -381,7 +384,8 @@ impl SmallVec { // SAFETY: all the members in 0..len are initialized let mut vec = Self { len: TaggedLen::new(len, false), - raw: RawSmallVec::new_inline(MaybeUninit::new(buf), Global) + raw: RawSmallVec::new_inline(MaybeUninit::new(buf)), + allocator: Global }; // Deallocate the remaining elements so no memory is leaked. unsafe { @@ -427,7 +431,8 @@ impl SmallVec { debug_assert!(len <= N); Self { len: TaggedLen::new(len, false), - raw: RawSmallVec::new_inline(buf, Global) + raw: RawSmallVec::new_inline(buf), + allocator: Global } } @@ -452,7 +457,8 @@ impl SmallVec { unsafe { vec.set_len(0) }; Self { len: TaggedLen::new(len, false), - raw: RawSmallVec::new(Global) + raw: RawSmallVec::new(), + allocator: Global } } else { let mut vec = ManuallyDrop::new(vec); @@ -464,7 +470,8 @@ impl SmallVec { Self { len: TaggedLen::new(len, true), - raw: RawSmallVec::new_heap(ptr, cap, Global) + raw: RawSmallVec::new_heap(ptr, cap), + allocator: Global } } } @@ -556,7 +563,8 @@ impl SmallVec { SmallVec { len: TaggedLen::new(length, true), - raw: RawSmallVec::new_heap(ptr, capacity, Global) + raw: RawSmallVec::new_heap(ptr, capacity), + allocator: Global } } } @@ -623,7 +631,7 @@ impl SmallVec { #[inline] pub const fn inline_size() -> usize { - RawSmallVec::::INLINE_CAP + RawSmallVec::::INLINE_CAP } #[inline] @@ -885,7 +893,10 @@ impl SmallVec { if new_capacity > Self::inline_size() { // SAFETY: we checked all the preconditions - let result = unsafe { self.raw.try_grow_raw(self.len, new_capacity) }; + let result = unsafe { + self.raw + .try_grow_raw(self.len, new_capacity, &self.allocator) + }; if result.is_ok() { // SAFETY: the allocation succeeded, so self.raw.heap is now @@ -898,7 +909,7 @@ impl SmallVec { if on_heap { unsafe { // SAFETY: heap member is active - let (ptr, old_cap) = self.raw.inner.heap; + let (ptr, old_cap) = self.raw.heap; // inline member is now active // SAFETY: len <= new_capacity <= Self::inline_size() @@ -908,7 +919,7 @@ impl SmallVec { ptr: ptr.cast(), size_bytes: old_cap * size_of::(), align: align_of::(), - alloc: &self.raw.alloc + alloc: &self.allocator }); self.set_inline(); } @@ -980,10 +991,10 @@ impl SmallVec { if len <= Self::inline_size() { // SAFETY: on_heap is true, so we're on the heap unsafe { - let (ptr, capacity) = self.raw.inner.heap; + let (ptr, capacity) = self.raw.heap; copy_nonoverlapping(ptr.as_ptr(), self.raw.as_mut_ptr_inline(), len); self.set_inline(); - self.raw.alloc.deallocate( + self.allocator.deallocate( ptr.cast(), Layout::from_size_align_unchecked(capacity * size_of::(), align_of::()) ); @@ -994,7 +1005,7 @@ impl SmallVec { // T can't be a ZST because SmallVec is never spilled. unsafe { self.raw - .try_grow_raw(self.len, len) + .try_grow_raw(self.len, len, &self.allocator) .unwrap_or_else(SmallVecError::handle) }; } @@ -1007,16 +1018,16 @@ impl SmallVec { return; } // SAFETY: the vector is on the heap - let capacity = unsafe { self.raw.inner.heap.1 }; + let capacity = unsafe { self.raw.heap.1 }; if capacity > min_capacity { let target = core::cmp::max(len, min_capacity); if target <= Self::inline_size() { // SAFETY: on_heap is true, so we're on the heap unsafe { - let (ptr, capacity) = self.raw.inner.heap; + let (ptr, capacity) = self.raw.heap; copy_nonoverlapping(ptr.as_ptr(), self.raw.as_mut_ptr_inline(), len); self.set_inline(); - self.raw.alloc.deallocate( + self.allocator.deallocate( ptr.cast(), Layout::from_size_align_unchecked( capacity * size_of::(), @@ -1030,7 +1041,7 @@ impl SmallVec { // T can't be a ZST because SmallVec is never spilled. unsafe { self.raw - .try_grow_raw(self.len, target) + .try_grow_raw(self.len, target, &self.allocator) .unwrap_or_else(SmallVecError::handle) }; } @@ -1226,7 +1237,7 @@ impl SmallVec { // - the first `len` entries are proper `T`-values // - the allocation is not larger than `isize::MAX` unsafe { - let (ptr, cap) = this.raw.inner.heap; + let (ptr, cap) = this.raw.heap; Vec::from_raw_parts(ptr.as_ptr(), len, cap) } } @@ -1577,7 +1588,8 @@ impl SmallVec { pub const fn new_in(alloc: A) -> SmallVec { Self { len: TaggedLen::new(0, false), - raw: RawSmallVec::new(alloc) + raw: RawSmallVec::new(), + allocator: alloc } } @@ -1585,7 +1597,10 @@ impl SmallVec { let mut this = Self::new_in(alloc); if capacity > Self::inline_size() && !Self::IS_ZST { // SAFETY: we checked all the preconditions - unsafe { this.raw.try_grow_raw(TaggedLen::new(0, false), capacity) }?; + unsafe { + this.raw + .try_grow_raw(TaggedLen::new(0, false), capacity, &this.allocator) + }?; // SAFETY: the allocation succeeded, so self.raw.heap is now active unsafe { this.set_on_heap() }; @@ -1671,7 +1686,7 @@ impl SmallVec { assert!(at <= len); let other_len = len - at; - let mut other = Self::with_capacity_in(other_len, self.raw.alloc.clone()); + let mut other = Self::with_capacity_in(other_len, self.allocator.clone()); // Unsafely `set_len` and copy items to `other`. unsafe { @@ -1726,12 +1741,12 @@ impl Drop for SmallVec { // SAFETY: see above unsafe { let _drop_dealloc = if on_heap { - let capacity = self.raw.inner.heap.1; + let capacity = self.raw.heap.1; Some(DropDealloc { ptr: NonNull::new_unchecked(ptr as *mut u8), size_bytes: capacity * size_of::(), align: align_of::(), - alloc: &self.raw.alloc + alloc: &self.allocator }) } else { None @@ -1749,12 +1764,12 @@ impl Drop for IntoIter { let begin = self.begin; let ptr = self.raw.as_mut_ptr(on_heap); let _drop_dealloc = if on_heap { - let capacity = self.raw.inner.heap.1; + let capacity = self.raw.heap.1; Some(DropDealloc { ptr: NonNull::new_unchecked(ptr as *mut u8), size_bytes: capacity * size_of::(), align: align_of::(), - alloc: &self.raw.alloc + alloc: &self.allocator }) } else { None @@ -1962,7 +1977,8 @@ impl Clone for SmallVec fn clone(&self) -> SmallVec { let mut vec = SmallVec { len: TaggedLen::new(0, false), - raw: RawSmallVec::new(self.raw.alloc.clone()) + raw: RawSmallVec::new(), + allocator: self.allocator.clone() }; vec.extend(self); @@ -1989,7 +2005,8 @@ impl Clone for IntoIter fn clone(&self) -> IntoIter { let mut vec = SmallVec { len: TaggedLen::new(0, false), - raw: RawSmallVec::new(self.raw.alloc.clone()) + raw: RawSmallVec::new(), + allocator: self.allocator.clone() }; vec.extend(self.as_slice()); @@ -2055,7 +2072,8 @@ impl IntoIterator for SmallVec { // the elements let this = ManuallyDrop::new(self); IntoIter { - raw: (&this.raw as *const RawSmallVec).read(), + raw: (&raw const this.raw).read(), + allocator: (&raw const this.allocator).read(), begin: 0, end: this.len } diff --git a/src/rawsmallvec.rs b/src/rawsmallvec.rs index c70d9aa2..ace0e6ad 100644 --- a/src/rawsmallvec.rs +++ b/src/rawsmallvec.rs @@ -18,47 +18,31 @@ use { }; #[repr(C)] -pub(crate) union RawSmallVecInner { - pub(crate) inline: ManuallyDrop>, - pub(crate) heap: (NonNull, usize) +pub union RawSmallVec { + pub inline: ManuallyDrop>, + pub heap: (NonNull, usize) } -/// Either a stack array with `length <= N` or a heap array -/// whose pointer and capacity are stored here. -/// -/// We store a `NonNull` instead of a `*mut T` so that type is covariant -/// with respect to `T`, and since the heap pointer is never null. -pub struct RawSmallVec { - pub(crate) inner: RawSmallVecInner, - pub(crate) alloc: A -} - -impl RawSmallVec { +impl RawSmallVec { pub const INLINE_CAP: usize = if Self::IS_ZST { usize::MAX } else { N }; const IS_ZST: bool = size_of::() == 0; #[inline] - pub const fn new(alloc: A) -> Self { - Self::new_inline(MaybeUninit::uninit(), alloc) + pub const fn new() -> Self { + Self::new_inline(MaybeUninit::uninit()) } #[inline] - pub const fn new_inline(inline: MaybeUninit<[T; N]>, alloc: A) -> Self { + pub const fn new_inline(inline: MaybeUninit<[T; N]>) -> Self { Self { - inner: RawSmallVecInner { - inline: ManuallyDrop::new(inline) - }, - alloc + inline: ManuallyDrop::new(inline) } } #[inline] - pub const fn new_heap(ptr: NonNull, capacity: usize, alloc: A) -> Self { + pub const fn new_heap(ptr: NonNull, capacity: usize) -> Self { Self { - inner: RawSmallVecInner { - heap: (ptr, capacity) - }, - alloc + heap: (ptr, capacity) } } @@ -68,14 +52,14 @@ impl RawSmallVec { // a reference to it. reading it would be UB potentially, but // for that downstream unsafe is required #[allow(unused_unsafe, reason = "Unsafe in MSRV")] - (unsafe { &raw const self.inner.inline }).cast() + (unsafe { &raw const self.inline }).cast() } #[inline] pub const fn as_mut_ptr_inline(&mut self) -> *mut T { // SAFETY: same as above #[allow(unused_unsafe, reason = "Unsafe in MSRV")] - (unsafe { &raw mut self.inner.inline }).cast() + (unsafe { &raw mut self.inline }).cast() } /// # Safety @@ -84,7 +68,7 @@ impl RawSmallVec { #[inline(always)] pub const unsafe fn as_ptr(&self, on_heap: bool) -> *const T { if on_heap { - unsafe { self.inner.heap.0.as_ptr() } + unsafe { self.heap.0.as_ptr() } } else { self.as_ptr_inline() } @@ -96,7 +80,7 @@ impl RawSmallVec { #[inline(always)] pub const unsafe fn as_mut_ptr(&mut self, on_heap: bool) -> *mut T { if on_heap { - unsafe { self.inner.heap.0.as_ptr() } + unsafe { self.heap.0.as_ptr() } } else { self.as_mut_ptr_inline() } @@ -108,7 +92,7 @@ impl RawSmallVec { #[inline(always)] pub const unsafe fn capacity(&self, on_heap: bool) -> usize { if on_heap { - unsafe { self.inner.heap.1 } + unsafe { self.heap.1 } } else { Self::INLINE_CAP } @@ -118,10 +102,13 @@ impl RawSmallVec { /// /// `new_capacity` must be non zero, and greater or equal to the length. /// T must not be a ZST. - pub unsafe fn try_grow_raw( + /// + /// the allocator must be the same one the data was allocated with + pub unsafe fn try_grow_raw( &mut self, len: TaggedLen, - new_capacity: usize + new_capacity: usize, + allocator: &A ) -> Result<(), SmallVecError> { let (len, was_on_heap) = len.parts(); debug_assert!(!Self::IS_ZST); @@ -139,8 +126,7 @@ impl RawSmallVec { let new_ptr = if !was_on_heap { // get a fresh allocation // `new_layout` has nonzero size. - let new_ptr = self - .alloc + let new_ptr = allocator .allocate(new_layout) .map_err(|_| SmallVecError::AllocationError(new_layout))? .cast(); @@ -152,10 +138,7 @@ impl RawSmallVec { // this can't overflow since we already constructed an equivalent // layout during the previous allocation let old_layout = unsafe { - Layout::from_size_align_unchecked( - self.inner.heap.1 * size_of::(), - align_of::() - ) + Layout::from_size_align_unchecked(self.heap.1 * size_of::(), align_of::()) }; // SAFETY: ptr was allocated with this allocator @@ -165,12 +148,12 @@ impl RawSmallVec { // alignment. since it was constructed // with Layout::array unsafe { - (if self.inner.heap.1 < new_capacity { + (if self.heap.1 < new_capacity { A::grow } else { A::shrink })( - &self.alloc, + allocator, NonNull::new(ptr as *mut u8).unwrap(), old_layout, new_layout @@ -179,7 +162,13 @@ impl RawSmallVec { .map_err(|_| SmallVecError::AllocationError(new_layout))? .cast() }; - self.inner.heap = (new_ptr, new_capacity); + self.heap = (new_ptr, new_capacity); Ok(()) } } + +impl Default for RawSmallVec { + fn default() -> Self { + Self::new() + } +}