Skip to content

Commit c34e154

Browse files
TennyZhuangoxalica
andcommitted
Optimize unnecessary check in Vec::retain
Co-authored-by: oxalica <[email protected]>
1 parent 4e1387c commit c34e154

File tree

1 file changed

+24
-3
lines changed

1 file changed

+24
-3
lines changed

library/alloc/src/vec/mod.rs

+24-3
Original file line numberDiff line numberDiff line change
@@ -1485,7 +1485,15 @@ impl<T, A: Allocator> Vec<T, A> {
14851485

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

1488-
while g.processed_len < original_len {
1488+
// process_one return a bool indicates whether the processing element should be retained.
1489+
#[inline(always)]
1490+
fn process_one<F, T, A: Allocator, const DELETED: bool>(
1491+
f: &mut F,
1492+
g: &mut BackshiftOnDrop<'_, T, A>,
1493+
) -> bool
1494+
where
1495+
F: FnMut(&T) -> bool,
1496+
{
14891497
// SAFETY: Unchecked element must be valid.
14901498
let cur = unsafe { &mut *g.v.as_mut_ptr().add(g.processed_len) };
14911499
if !f(cur) {
@@ -1495,9 +1503,9 @@ impl<T, A: Allocator> Vec<T, A> {
14951503
// SAFETY: We never touch this element again after dropped.
14961504
unsafe { ptr::drop_in_place(cur) };
14971505
// We already advanced the counter.
1498-
continue;
1506+
return false;
14991507
}
1500-
if g.deleted_cnt > 0 {
1508+
if DELETED {
15011509
// SAFETY: `deleted_cnt` > 0, so the hole slot must not overlap with current element.
15021510
// We use copy for move, and never touch this element again.
15031511
unsafe {
@@ -1506,6 +1514,19 @@ impl<T, A: Allocator> Vec<T, A> {
15061514
}
15071515
}
15081516
g.processed_len += 1;
1517+
return true;
1518+
}
1519+
1520+
// Stage 1: Nothing was deleted.
1521+
while g.processed_len < original_len {
1522+
if !process_one::<F, T, A, false>(&mut f, &mut g) {
1523+
break;
1524+
}
1525+
}
1526+
1527+
// Stage 2: Some elements were deleted.
1528+
while g.processed_len < original_len {
1529+
process_one::<F, T, A, true>(&mut f, &mut g);
15091530
}
15101531

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

0 commit comments

Comments
 (0)