perf/fix: retain mut algorithm

src/lib.rs

@@ -133,7 +133,8 @@ use {
         ptr::{
             NonNull,
             copy,
-            copy_nonoverlapping
+            copy_nonoverlapping,
+            drop_in_place
         }
     }
 };
@@ -1753,28 +1754,104 @@ impl<T, const N: usize> SmallVec<T, N> {
 
     #[inline]
     pub fn retain_mut<F: FnMut(&mut T) -> bool>(&mut self, mut f: F) {
-        let len = self.len();
+        let original_len = self.len();
 
-        if len == 0 {
-            // return early as hint to llvm, like what std does
+        if original_len == 0 {
+            // Empty case: explicit return allows better optimization, vs
+            // letting compiler infer it
             return;
         }
 
-        let ptr = self.as_mut_ptr();
-        let mut write_idx = 0;
+        // Vec: [Kept, Kept, Hole, Hole, Hole, Hole, Unchecked, Unchecked]
+        //      |            ^- write                ^- read             |
+        //      |<-              original_len                          ->|
+        // Kept: Elements which predicate returns true on.
+        // Hole: Moved or dropped element slot.
+        // Unchecked: Unchecked valid elements.
+        //
+        // This drop guard will be invoked when predicate or `drop` of element
+        // panicked. It shifts unchecked elements to cover holes and
+        // `set_len` to the correct length. In cases when predicate and
+        // `drop` never panic, it will be optimized out.
+        struct PanicGuard<'a, T, const N: usize> {
+            v: &'a mut SmallVec<T, N>,
+            read: usize,
+            write: usize,
+            original_len: usize
+        }
+
+        impl<T, const N: usize> Drop for PanicGuard<'_, T, N> {
+            #[cold]
+            fn drop(&mut self) {
+                let remaining = self.original_len - self.read;
+                // SAFETY: Trailing unchecked items must be valid since we never
+                // touch them.
+                unsafe {
+                    let ptr = self.v.as_mut_ptr();
+                    copy(ptr.add(self.read), ptr.add(self.write), remaining);
+                }
+                // SAFETY: After filling holes, all items are in contiguous
+                // memory.
+                unsafe {
+                    self.v.set_len(self.write + remaining);
+                }
+            }
+        }
 
-        for read_idx in 0..len {
-            // SAFETY: all the pointers are in bounds
-            unsafe {
-                if f(&mut *ptr.add(read_idx)) {
-                    if write_idx < read_idx {
-                        core::ptr::swap(ptr.add(read_idx), ptr.add(write_idx));
-                    }
-                    write_idx += 1;
+        let mut read = 0;
+        loop {
+            // SAFETY: read < original_len
+            let cur = unsafe { self.get_unchecked_mut(read) };
+            if !f(cur) {
+                break;
+            }
+            read += 1;
+            if read == original_len {
+                // All elements are kept, return early.
+                return;
+            }
+        }
+
+        // Critical section starts here and at least one element is going to be
+        // removed. Advance `g.read` early to avoid double drop if
+        // `drop_in_place` panicked.
+        let mut g = PanicGuard {
+            v: self,
+            read: read + 1,
+            write: read,
+            original_len
+        };
+        // SAFETY: previous `read` is always less than original_len.
+        unsafe { drop_in_place(g.v.as_mut_ptr().add(read)) }
+
+        while g.read < g.original_len {
+            let ptr = g.v.as_mut_ptr();
+            // SAFETY: `read` is always less than original_len.
+            let cur = unsafe { &mut *ptr.add(g.read) };
+            if !f(cur) {
+                // Advance `read` early to avoid double drop if `drop_in_place`
+                // panicked.
+                g.read += 1;
+                // SAFETY: We never touch this element again after dropped.
+                unsafe { drop_in_place(cur) };
+            } else {
+                // SAFETY: `read` > `write`, so the slots don't overlap.
+                // We use copy for move, and never touch the source element
+                // again.
+                unsafe {
+                    let hole = ptr.add(g.write);
+                    copy_nonoverlapping(cur, hole, 1);
                 }
+                g.write += 1;
+                g.read += 1;
             }
         }
-        self.truncate(write_idx);
+
+        // We are leaving the critical section and no panic happened,
+        // Commit the length change and forget the guard.
+        // SAFETY: `write` is always less than or equal to original_len.
+        unsafe { g.v.set_len(g.write) };
+        core::mem::forget(g);
     }
 
     #[inline]