diff --git a/library/alloc/src/vec/extract_if.rs b/library/alloc/src/vec/extract_if.rs index a456d3d9e602d..68c4ee5f81e10 100644 --- a/library/alloc/src/vec/extract_if.rs +++ b/library/alloc/src/vec/extract_if.rs @@ -64,27 +64,29 @@ where type Item = T; fn next(&mut self) -> Option { - unsafe { - while self.idx < self.end { - let i = self.idx; - let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len); - let drained = (self.pred)(&mut v[i]); - // Update the index *after* the predicate is called. If the index - // is updated prior and the predicate panics, the element at this - // index would be leaked. - self.idx += 1; - if drained { - self.del += 1; - return Some(ptr::read(&v[i])); - } else if self.del > 0 { - let del = self.del; - let src: *const T = &v[i]; - let dst: *mut T = &mut v[i - del]; - ptr::copy_nonoverlapping(src, dst, 1); + while self.idx < self.end { + let i = self.idx; + // SAFETY: Unchecked element must be valid. + let cur = unsafe { &mut *self.vec.as_mut_ptr().add(i) }; + let drained = (self.pred)(cur); + // Update the index *after* the predicate is called. If the index + // is updated prior and the predicate panics, the element at this + // index would be leaked. + self.idx += 1; + if drained { + self.del += 1; + // SAFETY: We never touch this element again after returning it. + return Some(unsafe { ptr::read(cur) }); + } else if self.del > 0 { + // SAFETY: `self.del` > 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 = self.vec.as_mut_ptr().add(i - self.del); + ptr::copy_nonoverlapping(cur, hole_slot, 1); } } - None } + None } fn size_hint(&self) -> (usize, Option) { @@ -95,14 +97,18 @@ where #[stable(feature = "extract_if", since = "1.87.0")] impl Drop for ExtractIf<'_, T, F, A> { fn drop(&mut self) { - unsafe { - if self.idx < self.old_len && self.del > 0 { - let ptr = self.vec.as_mut_ptr(); - let src = ptr.add(self.idx); - let dst = src.sub(self.del); - let tail_len = self.old_len - self.idx; - src.copy_to(dst, tail_len); + if self.del > 0 { + // SAFETY: Trailing unchecked items must be valid since we never touch them. + unsafe { + ptr::copy( + self.vec.as_ptr().add(self.idx), + self.vec.as_mut_ptr().add(self.idx - self.del), + self.old_len - self.idx, + ); } + } + // SAFETY: After filling holes, all items are in contiguous memory. + unsafe { self.vec.set_len(self.old_len - self.del); } }