Skip to content

avoid violating slice::from_raw_parts safety contract in Vec::extract_if #141032

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open

avoid violating slice::from_raw_parts safety contract in Vec::extract_if #141032

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 31 additions & 25 deletions library/alloc/src/vec/extract_if.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,27 +64,29 @@ where
type Item = T;

fn next(&mut self) -> Option<T> {
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.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is just stating a safety precondition of the mutable reference creation, but it doesn't explain why the operation is sound. The important thing here is that the index must be within the length of the original vector since it is smaller than self.end, which itself is smaller than the vectors original length.

Also, it would be helpful to add a comment that explains why we can't just use self.vec.get_unchecked_mut here (we can't since self.vec's length was set to zero, so indexing by i would violate its safety precondition and trigger the UB-check).

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<usize>) {
Expand All @@ -95,14 +97,18 @@ where
#[stable(feature = "extract_if", since = "1.87.0")]
impl<T, F, A: Allocator> 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);
}
}
Expand Down
Loading