From 42029c27c9243d8a178a0d64464185547b3a7d6f Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Tue, 8 Sep 2026 00:05:57 +0000 Subject: [PATCH 1/4] Compact retained elements directly in retain and dedup_by --- benches/compaction.rs | 100 ++++++++++++++++++++++++ src/lib.rs | 176 ++++++++++++++++++++++++++++++++++++------ tests/compaction.rs | 168 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 421 insertions(+), 23 deletions(-) create mode 100644 benches/compaction.rs create mode 100644 tests/compaction.rs diff --git a/benches/compaction.rs b/benches/compaction.rs new file mode 100644 index 00000000..16789cc0 --- /dev/null +++ b/benches/compaction.rs @@ -0,0 +1,100 @@ +#![feature(test)] +extern crate test; +use smallvec::SmallVec; +use test::{black_box, Bencher}; +fn bench(b: &mut Bencher, n: usize, mode: usize, wide: bool, retain: bool) { + if wide { + run::<[u64; 16]>(b, n, mode, retain) + } else { + run::<[u64; 1]>(b, n, mode, retain) + } +} +trait Payload: Copy { + fn from_key(key: u64) -> Self; + fn key(&self) -> u64; +} +impl Payload for [u64; 1] { + fn from_key(key: u64) -> Self { + [key; 1] + } + fn key(&self) -> u64 { + self[0] + } +} +impl Payload for [u64; 16] { + fn from_key(key: u64) -> Self { + [key; 16] + } + fn key(&self) -> u64 { + self[0] + } +} +fn run(b: &mut Bencher, n: usize, mode: usize, retain: bool) { + let input: Vec<_> = (0..n) + .map(|i| { + T::from_key(if retain || mode == 0 { + i as u64 + } else if mode == 1 { + (i / 2) as u64 + } else { + 0 + }) + }) + .collect(); + let mut values = SmallVec::<[T; 16]>::from_slice(&input); + // libtest includes restoring the input; both revisions use the same setup. + b.iter(|| { + values.clear(); + values.extend_from_slice(black_box(&input)); + if retain { + values.retain(|x| match mode { + 0 => true, + 1 => x.key() % 2 == 0, + _ => false, + }); + } else { + values.dedup_by(|a, b| a.key() == b.key()); + } + black_box(values.as_slice()); + }); +} +macro_rules! case { + ($name:ident,$n:expr,$mode:expr,$wide:expr,$retain:expr) => { + #[bench] + fn $name(b: &mut Bencher) { + bench(b, $n, $mode, $wide, $retain) + } + }; +} +case!(dedup_all_kept_16, 16, 0, false, false); +case!(dedup_half_kept_16, 16, 1, false, false); +case!(dedup_one_or_none_16, 16, 2, false, false); +case!(dedup_all_kept_17, 17, 0, false, false); +case!(dedup_half_kept_17, 17, 1, false, false); +case!(dedup_one_or_none_17, 17, 2, false, false); +case!(dedup_all_kept_4096, 4096, 0, true, false); +case!(dedup_half_kept_4096, 4096, 1, true, false); +case!(dedup_one_or_none_4096, 4096, 2, true, false); +case!(retain_all_kept_16, 16, 0, false, true); +case!(retain_half_kept_16, 16, 1, false, true); +case!(retain_one_or_none_16, 16, 2, false, true); +case!(retain_all_kept_17, 17, 0, false, true); +case!(retain_half_kept_17, 17, 1, false, true); +case!(retain_one_or_none_17, 17, 2, false, true); +case!(retain_all_kept_4096, 4096, 0, true, true); +case!(retain_half_kept_4096, 4096, 1, true, true); +case!(retain_one_or_none_4096, 4096, 2, true, true); +case!(dedup_u64_half_4096, 4096, 1, false, false); +case!(retain_u64_half_4096, 4096, 1, false, true); + +#[bench] +fn vec_dedup_pairs_4096_control(b: &mut Bencher) { + let input: Vec<_> = (0..4096).map(|i| [(i / 2) as u64; 16]).collect(); + let mut values = input.clone(); + b.iter(|| { + values.clear(); + values.extend_from_slice(black_box(&input)); + values.dedup_by(|a, b| a[0] == b[0]); + black_box(values.as_slice()); + }); +} diff --git a/src/lib.rs b/src/lib.rs index 74080244..622762b2 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. @@ -1658,31 +1746,73 @@ impl SmallVec { where F: FnMut(&mut A::Item, &mut A::Item) -> bool, { - // See the implementation of Vec::dedup_by in the - // standard library for an explanation of this algorithm. let len = self.len(); if len <= 1 { return; } - let ptr = self.as_mut_ptr(); - let mut w: usize = 1; + // Leave the unique prefix in place. Once we find a duplicate, + // every retained element can move directly into an earlier hole. + let mut read = 1; + while read < len { + let ptr = self.as_mut_ptr(); + // SAFETY: Both indices are initialized and distinct. + if unsafe { same_bucket(&mut *ptr.add(read), &mut *ptr.add(read - 1)) } { + break; + } + read += 1; + } + if read == len { + return; + } + + struct FillGapOnDrop<'a, A: Array> { + vec: &'a mut SmallVec, + read: usize, + write: usize, + } - unsafe { - for r in 1..len { - let p_r = ptr.add(r); - let p_wm1 = ptr.add(w - 1); - if !same_bucket(&mut *p_r, &mut *p_wm1) { - if r != w { - let p_w = p_wm1.add(1); - mem::swap(&mut *p_r, &mut *p_w); - } - w += 1; + impl Drop for FillGapOnDrop<'_, A> { + fn drop(&mut self) { + let remaining = self.vec.len() - self.read; + // SAFETY: The prefix before write and tail from read are + // initialized. Move the tail over the holes, even if the + // predicate or an element's destructor panics. + unsafe { + let ptr = self.vec.as_mut_ptr(); + ptr::copy(ptr.add(self.read), ptr.add(self.write), remaining); + self.vec.set_len(self.write + remaining); } } } - self.truncate(w); + let mut gap = FillGapOnDrop { + vec: self, + read: read + 1, + write: read, + }; + // SAFETY: read is the first duplicate. Advance the guard before + // dropping it so unwinding cannot drop the same element again. + unsafe { ptr::drop_in_place(gap.vec.as_mut_ptr().add(read)) }; + + let ptr = gap.vec.as_mut_ptr(); + // SAFETY: write < read <= len and write >= 1. Each survivor is + // moved once into a hole; moved-from slots are never dropped. + unsafe { + while gap.read < len { + let current = ptr.add(gap.read); + if same_bucket(&mut *current, &mut *ptr.add(gap.write - 1)) { + gap.read += 1; + ptr::drop_in_place(current); + } else { + ptr::copy_nonoverlapping(current, ptr.add(gap.write), 1); + gap.write += 1; + gap.read += 1; + } + } + gap.vec.set_len(gap.write); + } + core::mem::forget(gap); } /// Removes consecutive elements that map to the same key. diff --git a/tests/compaction.rs b/tests/compaction.rs new file mode 100644 index 00000000..4d8b04d2 --- /dev/null +++ b/tests/compaction.rs @@ -0,0 +1,168 @@ +use smallvec::SmallVec; +use std::{ + cell::Cell, + panic::{catch_unwind, AssertUnwindSafe}, + rc::Rc, +}; +type V = SmallVec<[T; 16]>; + +#[test] +fn compaction_matches_vec() { + for len in [0, 1, 2, 15, 16, 17, 64].iter().copied() { + for group in [1, 2, 3, 100].iter().copied() { + let input: Vec<_> = (0..len).map(|x| (x / group, x)).collect(); + let mut expected = input.clone(); + let mut actual: V<_> = input.into_iter().collect(); + expected.dedup_by(|a, b| { + b.1 += 1; + a.0 == b.0 + }); + actual.dedup_by(|a, b| { + b.1 += 1; + a.0 == b.0 + }); + assert_eq!(actual.as_slice(), expected.as_slice()); + } + } +} + +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() +} + +#[test] +fn dedup_predicate_panic_preserves_unprocessed_tail() { + for len in [8, 32].iter().copied() { + for panic_call in 0..len - 1 { + let (input, drops) = tracked(len, None); + let mut actual: V<_> = input.into_iter().collect(); + let mut calls = 0; + assert!(catch_unwind(AssertUnwindSafe(|| actual.dedup_by(|a, b| { + let call = calls; + calls += 1; + assert_ne!(call, panic_call); + a.id / 2 == b.id / 2 + }))) + .is_err()); + let read = panic_call + 1; + let expected: Vec<_> = (0..read).step_by(2).chain(read..len).collect(); + assert_eq!(ids(&actual), expected); + drop(actual); + assert!(drops.iter().all(|x| x.get() == 1)); + } + } +} + +#[test] +fn dedup_destructor_panic_preserves_unprocessed_tail() { + for len in [8, 32].iter().copied() { + for panic_at in (1..len).step_by(2) { + let (input, drops) = tracked(len, Some(panic_at)); + let mut actual: V<_> = input.into_iter().collect(); + assert!(catch_unwind(AssertUnwindSafe( + || actual.dedup_by(|a, b| a.id / 2 == b.id / 2) + )) + .is_err()); + let expected: Vec<_> = (0..panic_at).step_by(2).chain(panic_at + 1..len).collect(); + assert_eq!(ids(&actual), expected); + drop(actual); + assert!(drops.iter().all(|x| x.get() == 1)); + } + } +} + +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 dedup_zst_drops_once() { + ZST_DROPS.with(|x| x.set(0)); + let mut values: V<_> = (0..32).map(|_| Zst).collect(); + values.dedup_by(|_, _| true); + assert_eq!(values.len(), 1); + ZST_DROPS.with(|x| assert_eq!(x.get(), 31)); + drop(values); + ZST_DROPS.with(|x| assert_eq!(x.get(), 32)); +} + +#[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)); +} From 88c6bfabd610f34f4f1f9d530f58fb8a74b9a807 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Tue, 8 Sep 2026 02:10:57 +0000 Subject: [PATCH 2/4] Limit compaction optimization to retain --- benches/compaction.rs | 100 ----------------------------- benches/retain.rs | 76 ++++++++++++++++++++++ src/lib.rs | 72 +++++---------------- tests/{compaction.rs => retain.rs} | 72 --------------------- 4 files changed, 91 insertions(+), 229 deletions(-) delete mode 100644 benches/compaction.rs create mode 100644 benches/retain.rs rename tests/{compaction.rs => retain.rs} (54%) diff --git a/benches/compaction.rs b/benches/compaction.rs deleted file mode 100644 index 16789cc0..00000000 --- a/benches/compaction.rs +++ /dev/null @@ -1,100 +0,0 @@ -#![feature(test)] -extern crate test; -use smallvec::SmallVec; -use test::{black_box, Bencher}; -fn bench(b: &mut Bencher, n: usize, mode: usize, wide: bool, retain: bool) { - if wide { - run::<[u64; 16]>(b, n, mode, retain) - } else { - run::<[u64; 1]>(b, n, mode, retain) - } -} -trait Payload: Copy { - fn from_key(key: u64) -> Self; - fn key(&self) -> u64; -} -impl Payload for [u64; 1] { - fn from_key(key: u64) -> Self { - [key; 1] - } - fn key(&self) -> u64 { - self[0] - } -} -impl Payload for [u64; 16] { - fn from_key(key: u64) -> Self { - [key; 16] - } - fn key(&self) -> u64 { - self[0] - } -} -fn run(b: &mut Bencher, n: usize, mode: usize, retain: bool) { - let input: Vec<_> = (0..n) - .map(|i| { - T::from_key(if retain || mode == 0 { - i as u64 - } else if mode == 1 { - (i / 2) as u64 - } else { - 0 - }) - }) - .collect(); - let mut values = SmallVec::<[T; 16]>::from_slice(&input); - // libtest includes restoring the input; both revisions use the same setup. - b.iter(|| { - values.clear(); - values.extend_from_slice(black_box(&input)); - if retain { - values.retain(|x| match mode { - 0 => true, - 1 => x.key() % 2 == 0, - _ => false, - }); - } else { - values.dedup_by(|a, b| a.key() == b.key()); - } - black_box(values.as_slice()); - }); -} -macro_rules! case { - ($name:ident,$n:expr,$mode:expr,$wide:expr,$retain:expr) => { - #[bench] - fn $name(b: &mut Bencher) { - bench(b, $n, $mode, $wide, $retain) - } - }; -} -case!(dedup_all_kept_16, 16, 0, false, false); -case!(dedup_half_kept_16, 16, 1, false, false); -case!(dedup_one_or_none_16, 16, 2, false, false); -case!(dedup_all_kept_17, 17, 0, false, false); -case!(dedup_half_kept_17, 17, 1, false, false); -case!(dedup_one_or_none_17, 17, 2, false, false); -case!(dedup_all_kept_4096, 4096, 0, true, false); -case!(dedup_half_kept_4096, 4096, 1, true, false); -case!(dedup_one_or_none_4096, 4096, 2, true, false); -case!(retain_all_kept_16, 16, 0, false, true); -case!(retain_half_kept_16, 16, 1, false, true); -case!(retain_one_or_none_16, 16, 2, false, true); -case!(retain_all_kept_17, 17, 0, false, true); -case!(retain_half_kept_17, 17, 1, false, true); -case!(retain_one_or_none_17, 17, 2, false, true); -case!(retain_all_kept_4096, 4096, 0, true, true); -case!(retain_half_kept_4096, 4096, 1, true, true); -case!(retain_one_or_none_4096, 4096, 2, true, true); -case!(dedup_u64_half_4096, 4096, 1, false, false); -case!(retain_u64_half_4096, 4096, 1, false, true); - -#[bench] -fn vec_dedup_pairs_4096_control(b: &mut Bencher) { - let input: Vec<_> = (0..4096).map(|i| [(i / 2) as u64; 16]).collect(); - let mut values = input.clone(); - b.iter(|| { - values.clear(); - values.extend_from_slice(black_box(&input)); - values.dedup_by(|a, b| a[0] == b[0]); - black_box(values.as_slice()); - }); -} diff --git a/benches/retain.rs b/benches/retain.rs new file mode 100644 index 00000000..73dce2df --- /dev/null +++ b/benches/retain.rs @@ -0,0 +1,76 @@ +#![feature(test)] +extern crate test; +use smallvec::SmallVec; +use test::{black_box, Bencher}; +fn bench(b: &mut Bencher, n: usize, mode: usize, wide: bool) { + if wide { + run::<[u64; 16]>(b, n, mode) + } else { + run::<[u64; 1]>(b, n, mode) + } +} +trait Payload: Copy { + fn from_key(key: u64) -> Self; + fn key(&self) -> u64; +} +impl Payload for [u64; 1] { + fn from_key(key: u64) -> Self { + [key; 1] + } + fn key(&self) -> u64 { + self[0] + } +} +impl Payload for [u64; 16] { + fn from_key(key: u64) -> Self { + [key; 16] + } + fn key(&self) -> u64 { + self[0] + } +} +fn run(b: &mut Bencher, n: usize, mode: usize) { + let input: Vec<_> = (0..n).map(|i| T::from_key(i as u64)).collect(); + let mut values = SmallVec::<[T; 16]>::from_slice(&input); + // libtest includes restoring the input; both revisions use the same setup. + b.iter(|| { + values.clear(); + values.extend_from_slice(black_box(&input)); + values.retain(|x| match mode { + 0 => true, + 1 => x.key() % 2 == 0, + _ => false, + }); + black_box(values.as_slice()); + }); +} +macro_rules! case { + ($name:ident,$n:expr,$mode:expr,$wide:expr) => { + #[bench] + fn $name(b: &mut Bencher) { + bench(b, $n, $mode, $wide) + } + }; +} +case!(retain_all_kept_16, 16, 0, false); +case!(retain_half_kept_16, 16, 1, false); +case!(retain_one_or_none_16, 16, 2, false); +case!(retain_all_kept_17, 17, 0, false); +case!(retain_half_kept_17, 17, 1, false); +case!(retain_one_or_none_17, 17, 2, false); +case!(retain_all_kept_4096, 4096, 0, true); +case!(retain_half_kept_4096, 4096, 1, true); +case!(retain_one_or_none_4096, 4096, 2, true); +case!(retain_u64_half_4096, 4096, 1, false); + +#[bench] +fn vec_retain_half_4096_control(b: &mut Bencher) { + let input: Vec<_> = (0..4096).map(|i| [i as u64; 16]).collect(); + let mut values = input.clone(); + b.iter(|| { + values.clear(); + values.extend_from_slice(black_box(&input)); + values.retain(|x| x[0] % 2 == 0); + black_box(values.as_slice()); + }); +} diff --git a/src/lib.rs b/src/lib.rs index 622762b2..45bea51b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1746,73 +1746,31 @@ impl SmallVec { where F: FnMut(&mut A::Item, &mut A::Item) -> bool, { + // See the implementation of Vec::dedup_by in the + // standard library for an explanation of this algorithm. let len = self.len(); if len <= 1 { return; } - // Leave the unique prefix in place. Once we find a duplicate, - // every retained element can move directly into an earlier hole. - let mut read = 1; - while read < len { - let ptr = self.as_mut_ptr(); - // SAFETY: Both indices are initialized and distinct. - if unsafe { same_bucket(&mut *ptr.add(read), &mut *ptr.add(read - 1)) } { - break; - } - read += 1; - } - if read == len { - return; - } + let ptr = self.as_mut_ptr(); + let mut w: usize = 1; - struct FillGapOnDrop<'a, A: Array> { - vec: &'a mut SmallVec, - read: usize, - write: usize, - } - - impl Drop for FillGapOnDrop<'_, A> { - fn drop(&mut self) { - let remaining = self.vec.len() - self.read; - // SAFETY: The prefix before write and tail from read are - // initialized. Move the tail over the holes, even if the - // predicate or an element's destructor panics. - unsafe { - let ptr = self.vec.as_mut_ptr(); - ptr::copy(ptr.add(self.read), ptr.add(self.write), remaining); - self.vec.set_len(self.write + remaining); - } - } - } - - let mut gap = FillGapOnDrop { - vec: self, - read: read + 1, - write: read, - }; - // SAFETY: read is the first duplicate. Advance the guard before - // dropping it so unwinding cannot drop the same element again. - unsafe { ptr::drop_in_place(gap.vec.as_mut_ptr().add(read)) }; - - let ptr = gap.vec.as_mut_ptr(); - // SAFETY: write < read <= len and write >= 1. Each survivor is - // moved once into a hole; moved-from slots are never dropped. unsafe { - while gap.read < len { - let current = ptr.add(gap.read); - if same_bucket(&mut *current, &mut *ptr.add(gap.write - 1)) { - gap.read += 1; - ptr::drop_in_place(current); - } else { - ptr::copy_nonoverlapping(current, ptr.add(gap.write), 1); - gap.write += 1; - gap.read += 1; + for r in 1..len { + let p_r = ptr.add(r); + let p_wm1 = ptr.add(w - 1); + if !same_bucket(&mut *p_r, &mut *p_wm1) { + if r != w { + let p_w = p_wm1.add(1); + mem::swap(&mut *p_r, &mut *p_w); + } + w += 1; } } - gap.vec.set_len(gap.write); } - core::mem::forget(gap); + + self.truncate(w); } /// Removes consecutive elements that map to the same key. diff --git a/tests/compaction.rs b/tests/retain.rs similarity index 54% rename from tests/compaction.rs rename to tests/retain.rs index 4d8b04d2..eeb61630 100644 --- a/tests/compaction.rs +++ b/tests/retain.rs @@ -6,26 +6,6 @@ use std::{ }; type V = SmallVec<[T; 16]>; -#[test] -fn compaction_matches_vec() { - for len in [0, 1, 2, 15, 16, 17, 64].iter().copied() { - for group in [1, 2, 3, 100].iter().copied() { - let input: Vec<_> = (0..len).map(|x| (x / group, x)).collect(); - let mut expected = input.clone(); - let mut actual: V<_> = input.into_iter().collect(); - expected.dedup_by(|a, b| { - b.1 += 1; - a.0 == b.0 - }); - actual.dedup_by(|a, b| { - b.1 += 1; - a.0 == b.0 - }); - assert_eq!(actual.as_slice(), expected.as_slice()); - } - } -} - struct Tracked { id: usize, drops: Rc>>, @@ -52,47 +32,6 @@ fn ids(values: &[Tracked]) -> Vec { values.iter().map(|v| v.id).collect() } -#[test] -fn dedup_predicate_panic_preserves_unprocessed_tail() { - for len in [8, 32].iter().copied() { - for panic_call in 0..len - 1 { - let (input, drops) = tracked(len, None); - let mut actual: V<_> = input.into_iter().collect(); - let mut calls = 0; - assert!(catch_unwind(AssertUnwindSafe(|| actual.dedup_by(|a, b| { - let call = calls; - calls += 1; - assert_ne!(call, panic_call); - a.id / 2 == b.id / 2 - }))) - .is_err()); - let read = panic_call + 1; - let expected: Vec<_> = (0..read).step_by(2).chain(read..len).collect(); - assert_eq!(ids(&actual), expected); - drop(actual); - assert!(drops.iter().all(|x| x.get() == 1)); - } - } -} - -#[test] -fn dedup_destructor_panic_preserves_unprocessed_tail() { - for len in [8, 32].iter().copied() { - for panic_at in (1..len).step_by(2) { - let (input, drops) = tracked(len, Some(panic_at)); - let mut actual: V<_> = input.into_iter().collect(); - assert!(catch_unwind(AssertUnwindSafe( - || actual.dedup_by(|a, b| a.id / 2 == b.id / 2) - )) - .is_err()); - let expected: Vec<_> = (0..panic_at).step_by(2).chain(panic_at + 1..len).collect(); - assert_eq!(ids(&actual), expected); - drop(actual); - assert!(drops.iter().all(|x| x.get() == 1)); - } - } -} - thread_local! { static ZST_DROPS: Cell = Cell::new(0); } struct Zst; impl Drop for Zst { @@ -100,17 +39,6 @@ impl Drop for Zst { ZST_DROPS.with(|x| x.set(x.get() + 1)); } } -#[test] -fn dedup_zst_drops_once() { - ZST_DROPS.with(|x| x.set(0)); - let mut values: V<_> = (0..32).map(|_| Zst).collect(); - values.dedup_by(|_, _| true); - assert_eq!(values.len(), 1); - ZST_DROPS.with(|x| assert_eq!(x.get(), 31)); - drop(values); - ZST_DROPS.with(|x| assert_eq!(x.get(), 32)); -} - #[test] fn retain_panic_preserves_unprocessed_tail() { for len in [8, 32].iter().copied() { From 8d936338dbcefa380b4dd596379b57582ed5d616 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Tue, 8 Sep 2026 11:22:54 +0000 Subject: [PATCH 3/4] Remove added retain benchmark harness --- benches/retain.rs | 76 ----------------------------------------------- 1 file changed, 76 deletions(-) delete mode 100644 benches/retain.rs diff --git a/benches/retain.rs b/benches/retain.rs deleted file mode 100644 index 73dce2df..00000000 --- a/benches/retain.rs +++ /dev/null @@ -1,76 +0,0 @@ -#![feature(test)] -extern crate test; -use smallvec::SmallVec; -use test::{black_box, Bencher}; -fn bench(b: &mut Bencher, n: usize, mode: usize, wide: bool) { - if wide { - run::<[u64; 16]>(b, n, mode) - } else { - run::<[u64; 1]>(b, n, mode) - } -} -trait Payload: Copy { - fn from_key(key: u64) -> Self; - fn key(&self) -> u64; -} -impl Payload for [u64; 1] { - fn from_key(key: u64) -> Self { - [key; 1] - } - fn key(&self) -> u64 { - self[0] - } -} -impl Payload for [u64; 16] { - fn from_key(key: u64) -> Self { - [key; 16] - } - fn key(&self) -> u64 { - self[0] - } -} -fn run(b: &mut Bencher, n: usize, mode: usize) { - let input: Vec<_> = (0..n).map(|i| T::from_key(i as u64)).collect(); - let mut values = SmallVec::<[T; 16]>::from_slice(&input); - // libtest includes restoring the input; both revisions use the same setup. - b.iter(|| { - values.clear(); - values.extend_from_slice(black_box(&input)); - values.retain(|x| match mode { - 0 => true, - 1 => x.key() % 2 == 0, - _ => false, - }); - black_box(values.as_slice()); - }); -} -macro_rules! case { - ($name:ident,$n:expr,$mode:expr,$wide:expr) => { - #[bench] - fn $name(b: &mut Bencher) { - bench(b, $n, $mode, $wide) - } - }; -} -case!(retain_all_kept_16, 16, 0, false); -case!(retain_half_kept_16, 16, 1, false); -case!(retain_one_or_none_16, 16, 2, false); -case!(retain_all_kept_17, 17, 0, false); -case!(retain_half_kept_17, 17, 1, false); -case!(retain_one_or_none_17, 17, 2, false); -case!(retain_all_kept_4096, 4096, 0, true); -case!(retain_half_kept_4096, 4096, 1, true); -case!(retain_one_or_none_4096, 4096, 2, true); -case!(retain_u64_half_4096, 4096, 1, false); - -#[bench] -fn vec_retain_half_4096_control(b: &mut Bencher) { - let input: Vec<_> = (0..4096).map(|i| [i as u64; 16]).collect(); - let mut values = input.clone(); - b.iter(|| { - values.clear(); - values.extend_from_slice(black_box(&input)); - values.retain(|x| x[0] % 2 == 0); - black_box(values.as_slice()); - }); -} From 954d599e41b074282857e97c44211f0f797e6597 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 9 Sep 2026 10:54:12 -0400 Subject: [PATCH 4/4] Move retain tests into the unit test module --- src/tests.rs | 100 ++++++++++++++++++++++++++++++++++++++++++++++++ tests/retain.rs | 96 ---------------------------------------------- 2 files changed, 100 insertions(+), 96 deletions(-) delete mode 100644 tests/retain.rs diff --git a/src/tests.rs b/src/tests.rs index e0118ccb..2ff13f40 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]); diff --git a/tests/retain.rs b/tests/retain.rs deleted file mode 100644 index eeb61630..00000000 --- a/tests/retain.rs +++ /dev/null @@ -1,96 +0,0 @@ -use smallvec::SmallVec; -use std::{ - cell::Cell, - panic::{catch_unwind, AssertUnwindSafe}, - rc::Rc, -}; -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)); -}