Skip to content

Commit 0cdb41f

Browse files
committed
Mirror optimization of std::Vec::retain (rust-lang/rust#81126)
1 parent 83bf216 commit 0cdb41f

File tree

1 file changed

+48
-12
lines changed

1 file changed

+48
-12
lines changed

src/arrayvec.rs

+48-12
Original file line numberDiff line numberDiff line change
@@ -468,22 +468,58 @@ impl<A: Array> ArrayVec<A> {
468468
pub fn retain<F>(&mut self, mut f: F)
469469
where F: FnMut(&mut A::Item) -> bool
470470
{
471-
let len = self.len();
472-
let mut del = 0;
473-
{
474-
let v = &mut **self;
475-
476-
for i in 0..len {
477-
if !f(&mut v[i]) {
478-
del += 1;
479-
} else if del > 0 {
480-
v.swap(i - del, i);
471+
// Check the implementation of
472+
// https://doc.rust-lang.org/std/vec/struct.Vec.html#method.retain
473+
// for safety arguments (especially regarding panics in f and when
474+
// dropping elements). Implementation closely mirrored here.
475+
476+
let original_len = self.len();
477+
unsafe { self.set_len(0) };
478+
479+
struct BackshiftOnDrop<'a, A: Array> {
480+
v: &'a mut ArrayVec<A>,
481+
processed_len: usize,
482+
deleted_cnt: usize,
483+
original_len: usize,
484+
}
485+
486+
impl<A: Array> Drop for BackshiftOnDrop<'_, A> {
487+
fn drop(&mut self) {
488+
if self.deleted_cnt > 0 {
489+
unsafe {
490+
ptr::copy(
491+
self.v.as_ptr().add(self.processed_len),
492+
self.v.as_mut_ptr().add(self.processed_len - self.deleted_cnt),
493+
self.original_len - self.processed_len
494+
);
495+
}
496+
}
497+
unsafe {
498+
self.v.set_len(self.original_len - self.deleted_cnt);
481499
}
482500
}
483501
}
484-
if del > 0 {
485-
self.drain(len - del..);
502+
503+
let mut g = BackshiftOnDrop { v: self, processed_len: 0, deleted_cnt: 0, original_len };
504+
505+
while g.processed_len < original_len {
506+
let cur = unsafe { &mut *g.v.as_mut_ptr().add(g.processed_len) };
507+
if !f(cur) {
508+
g.processed_len += 1;
509+
g.deleted_cnt += 1;
510+
unsafe { ptr::drop_in_place(cur) };
511+
continue;
512+
}
513+
if g.deleted_cnt > 0 {
514+
unsafe {
515+
let hole_slot = g.v.as_mut_ptr().add(g.processed_len - g.deleted_cnt);
516+
ptr::copy_nonoverlapping(cur, hole_slot, 1);
517+
}
518+
}
519+
g.processed_len += 1;
486520
}
521+
522+
drop(g);
487523
}
488524

489525
/// Set the vector’s length without dropping or moving out elements

0 commit comments

Comments
 (0)