diff --git a/library/alloc/benches/vec.rs b/library/alloc/benches/vec.rs index 8e1d374b5d42b..0da4886278e39 100644 --- a/library/alloc/benches/vec.rs +++ b/library/alloc/benches/vec.rs @@ -733,11 +733,26 @@ fn bench_flat_map_collect(b: &mut Bencher) { b.iter(|| v.iter().flat_map(|color| color.rotate_left(8).to_be_bytes()).collect::>()); } +/// Reference benchmark that `retain` has to compete with. +#[bench] +fn bench_retain_iter_100000(b: &mut Bencher) { + let mut v = Vec::with_capacity(100000); + + b.iter(|| { + let mut tmp = std::mem::take(&mut v); + tmp.clear(); + tmp.extend(black_box(1..=100000)); + v = tmp.into_iter().filter(|x| x & 1 == 0).collect(); + }); +} + #[bench] fn bench_retain_100000(b: &mut Bencher) { - let v = (1..=100000).collect::>(); + let mut v = Vec::with_capacity(100000); + b.iter(|| { - let mut v = v.clone(); + v.clear(); + v.extend(black_box(1..=100000)); v.retain(|x| x & 1 == 0) }); } diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index 88bde6e8ce481..349e7105a92a4 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -1520,49 +1520,46 @@ impl Vec { let mut g = BackshiftOnDrop { v: self, processed_len: 0, deleted_cnt: 0, original_len }; - // process_one return a bool indicates whether the processing element should be retained. - #[inline(always)] - fn process_one( + fn process_loop( + original_len: usize, f: &mut F, g: &mut BackshiftOnDrop<'_, T, A>, - ) -> bool - where + ) where F: FnMut(&mut T) -> bool, { - // SAFETY: Unchecked element must be valid. - let cur = unsafe { &mut *g.v.as_mut_ptr().add(g.processed_len) }; - if !f(cur) { - // Advance early to avoid double drop if `drop_in_place` panicked. - g.processed_len += 1; - g.deleted_cnt += 1; - // SAFETY: We never touch this element again after dropped. - unsafe { ptr::drop_in_place(cur) }; - // We already advanced the counter. - return false; - } - if DELETED { - // SAFETY: `deleted_cnt` > 0, so the hole slot must not overlap with current element. - // We use copy for move, and never touch this element again. - unsafe { - let hole_slot = g.v.as_mut_ptr().add(g.processed_len - g.deleted_cnt); - ptr::copy_nonoverlapping(cur, hole_slot, 1); + while g.processed_len != original_len { + // SAFETY: Unchecked element must be valid. + let cur = unsafe { &mut *g.v.as_mut_ptr().add(g.processed_len) }; + if !f(cur) { + // Advance early to avoid double drop if `drop_in_place` panicked. + g.processed_len += 1; + g.deleted_cnt += 1; + // SAFETY: We never touch this element again after dropped. + unsafe { ptr::drop_in_place(cur) }; + // We already advanced the counter. + if DELETED { + continue; + } else { + break; + } + } + if DELETED { + // SAFETY: `deleted_cnt` > 0, so the hole slot must not overlap with current element. + // We use copy for move, and never touch this element again. + unsafe { + let hole_slot = g.v.as_mut_ptr().add(g.processed_len - g.deleted_cnt); + ptr::copy_nonoverlapping(cur, hole_slot, 1); + } } + g.processed_len += 1; } - g.processed_len += 1; - return true; } // Stage 1: Nothing was deleted. - while g.processed_len != original_len { - if !process_one::(&mut f, &mut g) { - break; - } - } + process_loop::(original_len, &mut f, &mut g); // Stage 2: Some elements were deleted. - while g.processed_len != original_len { - process_one::(&mut f, &mut g); - } + process_loop::(original_len, &mut f, &mut g); // All item are processed. This can be optimized to `set_len` by LLVM. drop(g);