Backport v2’s retain implementation to v1

src/lib.rs

@@ -1623,16 +1623,104 @@ impl<A: Array> SmallVec<A> {
     /// `false`. This method operates in place and preserves the order of
     /// the retained elements.
     pub fn retain<F: FnMut(&mut A::Item) -> bool>(&mut self, mut f: F) {
-        let mut del = 0;
-        let len = self.len();
-        for i in 0..len {
-            if !f(&mut self[i]) {
-                del += 1;
-            } else if del > 0 {
-                self.swap(i - del, i);
+        let original_len = self.len();
+
+        if original_len == 0 {
+            // Empty case: explicit return allows better optimization, vs
+            // letting compiler infer it
+            return;
+        }
+
+        // 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, A: Array> {
+            v: &'a mut SmallVec<A>,
+            read: usize,
+            write: usize,
+            original_len: usize,
+        }
+
+        impl<A: Array> Drop for PanicGuard<'_, A> {
+            #[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();
+                    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);
+                }
             }
         }
-        self.truncate(len - del);
+
+        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 { ptr::drop_in_place(g.v.as_mut_ptr().add(read)) }
+
+        let ptr = g.v.as_mut_ptr();
+        while g.read < g.original_len {
+            // 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 { ptr::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);
+                    ptr::copy_nonoverlapping(cur, hole, 1);
+                }
+                g.write += 1;
+                g.read += 1;
+            }
+        }
+
+        // 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);
     }
 
     /// Retains only the elements specified by the predicate.

src/tests.rs

@@ -786,6 +786,106 @@ fn test_retain() {
     assert_eq!(Rc::strong_count(&one), 1);
 }
 
+mod retain {
+    use crate::SmallVec;
+    use alloc::{rc::Rc, vec::Vec};
+    use std::{
+        cell::Cell,
+        panic::{catch_unwind, AssertUnwindSafe},
+        thread_local,
+    };
+    type V<T> = SmallVec<[T; 16]>;
+
+    struct Tracked {
+        id: usize,
+        drops: Rc<Vec<Cell<usize>>>,
+        panic_at: Option<usize>,
+    }
+    impl Drop for Tracked {
+        fn drop(&mut self) {
+            self.drops[self.id].set(self.drops[self.id].get() + 1);
+            assert_ne!(self.panic_at, Some(self.id), "drop panic");
+        }
+    }
+    fn tracked(len: usize, panic_at: Option<usize>) -> (Vec<Tracked>, Rc<Vec<Cell<usize>>>) {
+        let drops = Rc::new((0..len).map(|_| Cell::new(0)).collect::<Vec<_>>());
+        let values = (0..len)
+            .map(|id| Tracked {
+                id,
+                drops: drops.clone(),
+                panic_at,
+            })
+            .collect();
+        (values, drops)
+    }
+    fn ids(values: &[Tracked]) -> Vec<usize> {
+        values.iter().map(|v| v.id).collect()
+    }
+
+    thread_local! { static ZST_DROPS: Cell<usize> = Cell::new(0); }
+    struct Zst;
+    impl Drop for Zst {
+        fn drop(&mut self) {
+            ZST_DROPS.with(|x| x.set(x.get() + 1));
+        }
+    }
+    #[test]
+    fn retain_panic_preserves_unprocessed_tail() {
+        for len in [8, 32].iter().copied() {
+            for panic_at in 0..len {
+                for drop_panics in [false, true].iter().copied() {
+                    if drop_panics && panic_at % 2 == 0 {
+                        continue;
+                    }
+                    let destructor = if drop_panics { Some(panic_at) } else { None };
+                    let (input, drops) = tracked(len, destructor);
+                    let mut actual: V<_> = input.into_iter().collect();
+                    assert!(catch_unwind(AssertUnwindSafe(|| actual.retain(|x| {
+                        if !drop_panics {
+                            assert_ne!(x.id, panic_at);
+                        }
+                        x.id % 2 == 0
+                    })))
+                    .is_err());
+                    let read = panic_at + if drop_panics { 1 } else { 0 };
+                    let expected: Vec<_> = (0..panic_at).step_by(2).chain(read..len).collect();
+                    assert_eq!(ids(&actual), expected);
+                    drop(actual);
+                    assert!(drops.iter().all(|x| x.get() == 1));
+                }
+            }
+        }
+    }
+    #[test]
+    fn retain_patterns_and_zst() {
+        for len in [0, 1, 15, 16, 17, 64].iter().copied() {
+            for keep in 0..3 {
+                let mut actual: V<_> = (0..len).collect();
+                let mut expected: Vec<_> = (0..len).collect();
+                actual.retain(|x| {
+                    *x += 1;
+                    *x % 2 < keep
+                });
+                for x in &mut expected {
+                    *x += 1;
+                }
+                expected.retain(|x| *x % 2 < keep);
+                assert_eq!(actual.as_slice(), expected.as_slice());
+            }
+        }
+        ZST_DROPS.with(|x| x.set(0));
+        let mut values: V<_> = (0..32).map(|_| Zst).collect();
+        let mut seen = 0;
+        values.retain(|_| {
+            seen += 1;
+            seen % 2 == 0
+        });
+        assert_eq!(values.len(), 16);
+        drop(values);
+        ZST_DROPS.with(|x| assert_eq!(x.get(), 32));
+    }
+}
+
 #[test]
 fn test_dedup() {
     let mut dupes: SmallVec<[i32; 5]> = SmallVec::from_slice(&[1, 1, 2, 3, 3]);