mirror of
https://github.com/rust-lang/rust.git
synced 2025-06-05 11:48:30 +00:00
fix Drop items getting leaked in Filter::next_chunk
The optimization only makes sense for non-drop elements anyway. Use the default implementation for items that are Drop instead. It also simplifies the implementation.
This commit is contained in:
parent
c290e9de32
commit
133e7b10a4
@ -3,7 +3,7 @@ use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedF
|
|||||||
use crate::num::NonZero;
|
use crate::num::NonZero;
|
||||||
use crate::ops::Try;
|
use crate::ops::Try;
|
||||||
use core::array;
|
use core::array;
|
||||||
use core::mem::{ManuallyDrop, MaybeUninit};
|
use core::mem::MaybeUninit;
|
||||||
use core::ops::ControlFlow;
|
use core::ops::ControlFlow;
|
||||||
|
|
||||||
/// An iterator that filters the elements of `iter` with `predicate`.
|
/// An iterator that filters the elements of `iter` with `predicate`.
|
||||||
@ -27,6 +27,41 @@ impl<I, P> Filter<I, P> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl<I, P> Filter<I, P>
|
||||||
|
where
|
||||||
|
I: Iterator,
|
||||||
|
P: FnMut(&I::Item) -> bool,
|
||||||
|
{
|
||||||
|
#[inline]
|
||||||
|
fn next_chunk_dropless<const N: usize>(
|
||||||
|
&mut self,
|
||||||
|
) -> Result<[I::Item; N], array::IntoIter<I::Item, N>> {
|
||||||
|
let mut array: [MaybeUninit<I::Item>; N] = [const { MaybeUninit::uninit() }; N];
|
||||||
|
let mut initialized = 0;
|
||||||
|
|
||||||
|
let result = self.iter.try_for_each(|element| {
|
||||||
|
let idx = initialized;
|
||||||
|
initialized = idx + (self.predicate)(&element) as usize;
|
||||||
|
|
||||||
|
// SAFETY: Loop conditions ensure the index is in bounds.
|
||||||
|
unsafe { array.get_unchecked_mut(idx) }.write(element);
|
||||||
|
|
||||||
|
if initialized < N { ControlFlow::Continue(()) } else { ControlFlow::Break(()) }
|
||||||
|
});
|
||||||
|
|
||||||
|
match result {
|
||||||
|
ControlFlow::Break(()) => {
|
||||||
|
// SAFETY: The loop above is only explicitly broken when the array has been fully initialized
|
||||||
|
Ok(unsafe { MaybeUninit::array_assume_init(array) })
|
||||||
|
}
|
||||||
|
ControlFlow::Continue(()) => {
|
||||||
|
// SAFETY: The range is in bounds since the loop breaks when reaching N elements.
|
||||||
|
Err(unsafe { array::IntoIter::new_unchecked(array, 0..initialized) })
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[stable(feature = "core_impl_debug", since = "1.9.0")]
|
#[stable(feature = "core_impl_debug", since = "1.9.0")]
|
||||||
impl<I: fmt::Debug, P> fmt::Debug for Filter<I, P> {
|
impl<I: fmt::Debug, P> fmt::Debug for Filter<I, P> {
|
||||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||||
@ -64,52 +99,15 @@ where
|
|||||||
fn next_chunk<const N: usize>(
|
fn next_chunk<const N: usize>(
|
||||||
&mut self,
|
&mut self,
|
||||||
) -> Result<[Self::Item; N], array::IntoIter<Self::Item, N>> {
|
) -> Result<[Self::Item; N], array::IntoIter<Self::Item, N>> {
|
||||||
let mut array: [MaybeUninit<Self::Item>; N] = [const { MaybeUninit::uninit() }; N];
|
let fun = const {
|
||||||
|
if crate::mem::needs_drop::<I::Item>() {
|
||||||
struct Guard<'a, T> {
|
array::iter_next_chunk::<I::Item, N>
|
||||||
array: &'a mut [MaybeUninit<T>],
|
} else {
|
||||||
initialized: usize,
|
Self::next_chunk_dropless::<N>
|
||||||
}
|
|
||||||
|
|
||||||
impl<T> Drop for Guard<'_, T> {
|
|
||||||
#[inline]
|
|
||||||
fn drop(&mut self) {
|
|
||||||
if const { crate::mem::needs_drop::<T>() } {
|
|
||||||
// SAFETY: self.initialized is always <= N, which also is the length of the array.
|
|
||||||
unsafe {
|
|
||||||
core::ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(
|
|
||||||
self.array.get_unchecked_mut(..self.initialized),
|
|
||||||
));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
};
|
||||||
|
|
||||||
let mut guard = Guard { array: &mut array, initialized: 0 };
|
fun(self)
|
||||||
|
|
||||||
let result = self.iter.try_for_each(|element| {
|
|
||||||
let idx = guard.initialized;
|
|
||||||
guard.initialized = idx + (self.predicate)(&element) as usize;
|
|
||||||
|
|
||||||
// SAFETY: Loop conditions ensure the index is in bounds.
|
|
||||||
unsafe { guard.array.get_unchecked_mut(idx) }.write(element);
|
|
||||||
|
|
||||||
if guard.initialized < N { ControlFlow::Continue(()) } else { ControlFlow::Break(()) }
|
|
||||||
});
|
|
||||||
|
|
||||||
let guard = ManuallyDrop::new(guard);
|
|
||||||
|
|
||||||
match result {
|
|
||||||
ControlFlow::Break(()) => {
|
|
||||||
// SAFETY: The loop above is only explicitly broken when the array has been fully initialized
|
|
||||||
Ok(unsafe { MaybeUninit::array_assume_init(array) })
|
|
||||||
}
|
|
||||||
ControlFlow::Continue(()) => {
|
|
||||||
let initialized = guard.initialized;
|
|
||||||
// SAFETY: The range is in bounds since the loop breaks when reaching N elements.
|
|
||||||
Err(unsafe { array::IntoIter::new_unchecked(array, 0..initialized) })
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[inline]
|
#[inline]
|
||||||
|
Loading…
Reference in New Issue
Block a user