diff --git a/src/lib.rs b/src/lib.rs index f7b1bba..40abe8b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1610,21 +1610,7 @@ impl SmallVec { #[inline] pub fn retain_mut bool>(&mut self, mut f: F) { - let mut del = 0; - let len = self.len(); - let ptr = self.as_mut_ptr(); - for i in 0..len { - // SAFETY: all the pointers are in bounds - // `i - del` never overflows since `del <= i` is a maintained invariant - unsafe { - if !f(&mut *ptr.add(i)) { - del += 1; - } else if del > 0 { - core::ptr::swap(ptr.add(i), ptr.add(i - del)); - } - } - } - self.truncate(len - del); + self.extract_if(.., |element| !f(element)).for_each(drop); } #[inline] diff --git a/src/tests.rs b/src/tests.rs index cfd801b..a53ac4a 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -1,5 +1,6 @@ use crate::{smallvec, SmallVec}; +use core::cell::{Cell, RefCell}; use core::hash::Hasher; use core::iter::FromIterator; @@ -748,6 +749,111 @@ fn test_retain() { assert_eq!(Rc::strong_count(&one), 1); } +fn check_retain_mut_predicate_panic_matches_vec() { + struct Counted { + value: i32, + drops: Rc>, + } + + impl Drop for Counted { + fn drop(&mut self) { + self.drops.set(self.drops.get() + 1); + } + } + + let drops = Rc::new(Cell::new(0)); + let mut values = SmallVec::::new(); + values.extend((1..=4).map(|value| Counted { + value, + drops: Rc::clone(&drops), + })); + + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + values.retain_mut(|item| match item.value { + 1 => false, + 3 => panic!("predicate panic"), + _ => true, + }); + })); + + assert!(result.is_err()); + assert_eq!(drops.get(), 1); + assert_eq!( + values.iter().map(|item| item.value).collect::>(), + vec![2, 3, 4] + ); + + drop(values); + assert_eq!(drops.get(), 4); +} + +#[test] +fn test_retain_mut_predicate_panic_inline() { + check_retain_mut_predicate_panic_matches_vec::<4>(); +} + +#[test] +fn test_retain_mut_predicate_panic_spilled() { + check_retain_mut_predicate_panic_matches_vec::<2>(); +} + +fn check_retain_mut_drop_panic_stops_visiting() { + struct PanicOnDrop { + value: i32, + drops: Rc>>, + panic_once: Rc>, + } + + impl Drop for PanicOnDrop { + fn drop(&mut self) { + self.drops.borrow_mut().push(self.value); + if self.value == 1 && !self.panic_once.replace(true) { + panic!("drop panic"); + } + } + } + + let drops = Rc::new(RefCell::new(Vec::new())); + let visited = Rc::new(RefCell::new(Vec::new())); + let panic_once = Rc::new(Cell::new(false)); + let mut values = SmallVec::::new(); + values.extend((1..=4).map(|value| PanicOnDrop { + value, + drops: Rc::clone(&drops), + panic_once: Rc::clone(&panic_once), + })); + + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + values.retain_mut(|item| { + visited.borrow_mut().push(item.value); + item.value != 1 + }); + })); + + assert!(result.is_err()); + assert_eq!(&*visited.borrow(), &[1]); + assert_eq!(&*drops.borrow(), &[1]); + assert_eq!( + values.iter().map(|item| item.value).collect::>(), + vec![2, 3, 4] + ); + + drop(values); + let mut dropped = drops.borrow().clone(); + dropped.sort_unstable(); + assert_eq!(dropped, [1, 2, 3, 4]); +} + +#[test] +fn test_retain_mut_drop_panic_inline() { + check_retain_mut_drop_panic_stops_visiting::<4>(); +} + +#[test] +fn test_retain_mut_drop_panic_spilled() { + check_retain_mut_drop_panic_stops_visiting::<2>(); +} + #[test] fn test_dedup() { let mut dupes: SmallVec = SmallVec::from(&[1, 1, 2, 3, 3]);