Improved performance and solved performance regressions of v2 against v1

src/lib.rs

@@ -207,16 +207,12 @@ impl<T, const N: usize> RawSmallVec<T, N> {
 
     #[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<T, const N: usize> SmallVec<T, N> {
 
     #[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<T, const N: usize> SmallVec<T, N> {
     #[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<T: Clone, const N: usize> SmallVec<T, N> {
 
     #[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<R>(&mut self, src: R)
@@ -2403,9 +2450,37 @@ impl<T, const N: usize> SmallVec<T, N> {
     where
         I: IntoIterator<Item = T>,
     {
-        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<T, N> {
 
         #[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());
+            }
         }
     }
 }

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<usize>,
+        drops: &'a Cell<usize>,
+    }
+    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<PanicOnClone<'_>, 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<u8, 8> = smallvec![0, 1, 2, 3];