diff --git a/src/lib.rs b/src/lib.rs index 7408024..45bea51 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1623,16 +1623,104 @@ impl SmallVec { /// `false`. This method operates in place and preserves the order of /// the retained elements. pub fn retain 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, + read: usize, + write: usize, + original_len: usize, + } + + impl 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. diff --git a/src/tests.rs b/src/tests.rs index e0118cc..2ff13f4 100644 --- a/src/tests.rs +++ b/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 = SmallVec<[T; 16]>; + + struct Tracked { + id: usize, + drops: Rc>>, + panic_at: Option, + } + 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) -> (Vec, Rc>>) { + let drops = Rc::new((0..len).map(|_| Cell::new(0)).collect::>()); + let values = (0..len) + .map(|id| Tracked { + id, + drops: drops.clone(), + panic_at, + }) + .collect(); + (values, drops) + } + fn ids(values: &[Tracked]) -> Vec { + values.iter().map(|v| v.id).collect() + } + + thread_local! { static ZST_DROPS: Cell = 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]);