Skip to content

Commit

Permalink
Optimize unnecessary check in Vec::retain
Browse files Browse the repository at this point in the history
Co-authored-by: oxalica <[email protected]>
  • Loading branch information
TennyZhuang and oxalica committed Sep 16, 2021
1 parent 20e14e4 commit 3839ca9
Showing 1 changed file with 24 additions and 3 deletions.
27 changes: 24 additions & 3 deletions library/alloc/src/vec/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1485,7 +1485,15 @@ impl<T, A: Allocator> Vec<T, A> {

let mut g = BackshiftOnDrop { v: self, processed_len: 0, deleted_cnt: 0, original_len };

while g.processed_len < original_len {
// process_one return a bool indicates whether the processing element should be retained.
#[inline(always)]
fn process_one<F, T, A: Allocator, const DELETED: bool>(
f: &mut F,
g: &mut BackshiftOnDrop<'_, T, A>,
) -> bool
where
F: FnMut(&T) -> bool,
{
// SAFETY: Unchecked element must be valid.
let cur = unsafe { &mut *g.v.as_mut_ptr().add(g.processed_len) };
if !f(cur) {
Expand All @@ -1495,9 +1503,9 @@ impl<T, A: Allocator> Vec<T, A> {
// SAFETY: We never touch this element again after dropped.
unsafe { ptr::drop_in_place(cur) };
// We already advanced the counter.
continue;
return false;
}
if g.deleted_cnt > 0 {
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 {
Expand All @@ -1506,6 +1514,19 @@ impl<T, A: Allocator> Vec<T, A> {
}
}
g.processed_len += 1;
return true;
}

// Stage 1: Nothing was deleted.
while g.processed_len != original_len {
if !process_one::<F, T, A, false>(&mut f, &mut g) {
break;
}
}

// Stage 2: Some elements were deleted.
while g.processed_len != original_len {
process_one::<F, T, A, true>(&mut f, &mut g);
}

// All item are processed. This can be optimized to `set_len` by LLVM.
Expand Down

0 comments on commit 3839ca9

Please sign in to comment.