diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index d2bc5d9bce39d..2255e842cc5ce 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -2297,8 +2297,7 @@ impl Vec { unsafe { self.set_len(0) }; // Vec: [Kept, Kept, Hole, Hole, Hole, Hole, Unchecked, Unchecked] - // |<- processed len ->| ^- next to check - // |<- deleted cnt ->| + // | ^- write_index ^- read_index // |<- original_len ->| // Kept: Elements which predicate returns true on. // Hole: Moved or dropped element slot. @@ -2309,31 +2308,39 @@ impl Vec { // In cases when predicate and `drop` never panick, it will be optimized out. struct BackshiftOnDrop<'a, T, A: Allocator> { v: &'a mut Vec, - processed_len: usize, - deleted_cnt: usize, + read_index: usize, + write_index: usize, original_len: usize, } impl Drop for BackshiftOnDrop<'_, T, A> { fn drop(&mut self) { - if self.deleted_cnt > 0 { + if self.read_index < self.original_len { + self.write_index = self.write_index.min(self.read_index); + let remaining = self.original_len - self.read_index; // SAFETY: Trailing unchecked items must be valid since we never touch them. unsafe { ptr::copy( - self.v.as_ptr().add(self.processed_len), - self.v.as_mut_ptr().add(self.processed_len - self.deleted_cnt), - self.original_len - self.processed_len, + self.v.as_ptr().add(self.read_index), + self.v.as_mut_ptr().add(self.write_index), + remaining, ); } + // SAFETY: After filling holes, all items are in contiguous memory. + unsafe { + self.v.set_len(self.write_index + remaining); + } + return; } - // SAFETY: After filling holes, all items are in contiguous memory. + // SAFETY: no panic happened, so all items are in contiguous memory. unsafe { - self.v.set_len(self.original_len - self.deleted_cnt); + self.v.set_len(self.write_index.min(self.read_index)); } } } - let mut g = BackshiftOnDrop { v: self, processed_len: 0, deleted_cnt: 0, original_len }; + let mut g = + BackshiftOnDrop { v: self, read_index: 0, write_index: usize::MAX, original_len }; fn process_loop( original_len: usize, @@ -2342,31 +2349,37 @@ impl Vec { ) where F: FnMut(&mut T) -> bool, { - while g.processed_len != original_len { + while g.read_index < original_len { // SAFETY: Unchecked element must be valid. - let cur = unsafe { &mut *g.v.as_mut_ptr().add(g.processed_len) }; + let cur = unsafe { &mut *g.v.as_mut_ptr().add(g.read_index) }; if !f(cur) { - // Advance early to avoid double drop if `drop_in_place` panicked. - g.processed_len += 1; - g.deleted_cnt += 1; + if !DELETED { + // We set g.write_index only once when we found the first deleted element. + g.write_index = g.read_index; + } + // Advance read_index early to avoid double drop if `drop_in_place` panicked. + g.read_index += 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 { + if !DELETED { + // We found the first deleted element. + // Switch to the second stage. + // read_index is now always greater than write_index. 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); + } else { + if DELETED { + // SAFETY: `read_index` > `write_index`, so the slots don't overlap. + // We use copy for move, and never touch the source element again. + unsafe { + let hole_slot = g.v.as_mut_ptr().add(g.write_index); + ptr::copy_nonoverlapping(cur, hole_slot, 1); + } + g.write_index += 1; } + g.read_index += 1; } - g.processed_len += 1; } } @@ -2376,7 +2389,7 @@ impl Vec { // Stage 2: Some elements were deleted. process_loop::(original_len, &mut f, &mut g); - // All item are processed. This can be optimized to `set_len` by LLVM. + // All items are processed. This can be optimized to `set_len` by LLVM. drop(g); }