diff --git a/src/lib.rs b/src/lib.rs index e6bc8c7..c2987d3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -207,16 +207,12 @@ impl RawSmallVec { #[inline] const fn as_ptr_inline(&self) -> *const T { - // SAFETY: This is safe because we don't read the value. We only get a pointer to the data. - // Dereferencing the pointer is unsafe so unsafe code is required to misuse the return - // value. - (unsafe { addr_of!(self.inline) }) as *const T + unsafe { addr_of!(self.inline) as *const T } } #[inline] const fn as_mut_ptr_inline(&mut self) -> *mut T { - // SAFETY: See above. - (unsafe { addr_of_mut!(self.inline) }) as *mut T + unsafe { addr_of_mut!(self.inline) as *mut T } } /// # Safety @@ -1223,16 +1219,37 @@ impl SmallVec { #[inline] pub fn push(&mut self, value: T) { - let len = self.len(); - if len == self.capacity() { + let on_heap = self.len.on_heap(); + let len = self.len.value(); + let cap = if on_heap { + // SAFETY: `on_heap` being true means the active union variant is `heap`. + unsafe { self.raw.heap.1 } + } else { + Self::inline_size() + }; + + if len < cap { + // SAFETY: both the input and output are within the allocation and we allocated enough + // space in case it wasn't enough, so the address is valid for writes + unsafe { + let ptr = if on_heap { + self.raw.as_mut_ptr_heap() + } else { + self.raw.as_mut_ptr_inline() + }; + ptr.add(len).write(value); + } + // saves one on_heap function call, it hasn't changed because we haven't resized + self.len = TaggedLen::new(len + 1, on_heap); + } else { self.reserve(1); + // SAFETY: both the input and output are within the allocation and we allocated enough + // space in case it wasn't enough, so the address is valid for writes + unsafe { + self.as_mut_ptr().add(len).write(value); + self.set_len(len + 1); + } } - // SAFETY: both the input and output are within the allocation - let ptr = unsafe { self.as_mut_ptr().add(len) }; - // SAFETY: we allocated enough space in case it wasn't enough, so the address is valid for - // writes. - unsafe { ptr.write(value) }; - unsafe { self.set_len(len + 1) } } #[inline] @@ -1498,19 +1515,25 @@ impl SmallVec { #[inline] pub fn insert(&mut self, index: usize, value: T) { let len = self.len(); - assert!(index <= len, "insertion index (is {index}) should be <= len (is {len})"); - self.reserve(1); - let ptr = self.as_mut_ptr(); - unsafe { - // the elements at `index + 1..len + 1` are now initialized - if index < len { - copy(ptr.add(index), ptr.add(index + 1), len - index); + if index == len { + // inserting at length is just pushing + self.push(value); + } else { + assert!(index < len, "insertion index (is {index}) should be <= len (is {len})"); + if len == self.capacity() { + self.reserve(1); } - // the element at `index` is now initialized - ptr.add(index).write(value); + let ptr = self.as_mut_ptr(); + unsafe { + // the elements at `index + 1..len + 1` are now initialized + copy(ptr.add(index), ptr.add(index + 1), len - index); + + // the element at `index` is now initialized + ptr.add(index).write(value); - // SAFETY: all the elements are initialized - self.set_len(len + 1); + // SAFETY: all the elements are initialized + self.set_len(len + 1); + } } } @@ -1799,7 +1822,31 @@ impl SmallVec { #[inline] pub fn extend_from_slice(&mut self, other: &[T]) { - self.extend(other.iter()) + let len = self.len.value(); + let other_length = other.len(); + let cap = self.capacity(); + // allocate needed capacity if needed, not more + if len + other_length > cap {self.reserve(other_length + len - cap);} + let ptr = self.as_mut_ptr(); + + // SAFETY: capacity has been reserved so other.len spots are guaranted + unsafe { + let mut guard = DropGuard { + // starts from len + ptr: ptr.add(len), + len: 0, + }; + for item in other { + ptr.add(len + guard.len).write(item.clone()); + guard.len += 1; + } + + // update length. guard.len = other.len + self.set_len(len + guard.len); + + // guard no longer needed, forgetting avoids double free + core::mem::forget(guard); + } } pub fn extend_from_within(&mut self, src: R) @@ -2403,9 +2450,37 @@ impl SmallVec { where I: IntoIterator, { - let iter = iter.into_iter(); + let mut iter = iter.into_iter(); let (size, _) = iter.size_hint(); self.reserve(size); + + let len = self.len(); + let cap = self.capacity(); + let ptr = self.as_mut_ptr(); + + // SAFETY: only operates as long as there's capacity, then it defaults to push + unsafe { + let mut guard = DropGuard { + ptr: ptr.add(len), + len: 0, + }; + + // while-let syntax is not allowed in 1.83 2021 + while len + guard.len < cap { + let Some(x) = iter.next() else { + self.set_len(len + guard.len); + core::mem::forget(guard); + return; + }; + ptr.add(len + guard.len).write(x); + guard.len += 1; + } + + self.set_len(len + guard.len); + core::mem::forget(guard); + } + + // after capacity has been exhausted, push items for x in iter { self.push(x); } @@ -2641,7 +2716,34 @@ impl<'a, T: Clone + 'a, const N: usize> Extend<&'a T> for SmallVec { #[cfg(not(feature = "specialization"))] { - self.extend_fallback(iter.into_iter().cloned()); + let mut iter = iter.into_iter(); + + // be conservative reserving size, reserve lower bound + let mut len = self.len.value(); + let size = iter.size_hint().0; + // could be further optimized to exactly reserve the amount of elements at upper or + // lower bound, but performance gains might be marginal knowing it's just an estimate + self.reserve(size); + + let on_heap = self.len.on_heap(); + let cap = self.capacity(); + let ptr = self.as_mut_ptr(); + + while len < cap { + let Some(x) = iter.next() else { + return; + }; + unsafe { + ptr.add(len).write(x.clone()); + } + len += 1; + self.len = TaggedLen::new(len, on_heap); + } + + // capacity exhausted, default to push + for x in iter { + self.push(x.clone()); + } } } } diff --git a/src/tests.rs b/src/tests.rs index 2082f96..0e706c8 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -387,6 +387,84 @@ fn test_extend_from_slice() { ); } +#[test] +fn test_extend_from_slice_clone_panic_safety() { + use core::cell::Cell; + use core::panic::AssertUnwindSafe; + // needs std to catch unwind + use std::panic::catch_unwind; + + // mimick panicking on clone + #[derive(Debug)] + struct PanicOnClone<'a> { + value: u8, + budget: &'a Cell, + drops: &'a Cell, + } + impl Clone for PanicOnClone<'_> { + fn clone(&self) -> Self { + let remaining = self.budget.get(); + if remaining == 0 { + panic!("clone panic"); + } + self.budget.set(remaining - 1); + Self { + value: self.value, + budget: self.budget, + drops: self.drops, + } + } + } + impl Drop for PanicOnClone<'_> { + fn drop(&mut self) { + self.drops.set(self.drops.get() + 1); + } + } + + // dies on purpose after two clones + let budget = Cell::new(2); + let drops = Cell::new(0); + let mut dst: SmallVec, 8> = SmallVec::new(); + dst.push(PanicOnClone { + value: 10, + budget: &budget, + drops: &drops, + }); + + // test with example data + let src = [ + PanicOnClone { + value: 11, + budget: &budget, + drops: &drops, + }, + PanicOnClone { + value: 12, + budget: &budget, + drops: &drops, + }, + PanicOnClone { + value: 13, + budget: &budget, + drops: &drops, + }, + ]; + + // check it panicked + let panicked = catch_unwind(AssertUnwindSafe(|| { + dst.extend_from_slice(&src); + })) + .is_err(); + assert!(panicked); + assert_eq!(dst.len(), 1); + assert_eq!(dst[0].value, 10); + + // check drop counters + drop(dst); + drop(src); + assert_eq!(drops.get(), 6); +} + #[test] fn test_extend_from_within() { let mut v: SmallVec = smallvec![0, 1, 2, 3];