From 133e7b10a45bddd4ef5ba265d848ce0a1976e8ae Mon Sep 17 00:00:00 2001
From: The 8472 <git@infinite-source.de>
Date: Sun, 23 Jun 2024 22:16:22 +0200
Subject: [PATCH 1/3] 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.
---
library/core/src/iter/adapters/filter.rs | 88 ++++++++++++------------
1 file changed, 43 insertions(+), 45 deletions(-)
diff --git a/library/core/src/iter/adapters/filter.rs b/library/core/src/iter/adapters/filter.rs
index ca23d1b13a88b..1c99f3938e270 100644
--- a/library/core/src/iter/adapters/filter.rs
+++ b/library/core/src/iter/adapters/filter.rs
@@ -3,7 +3,7 @@ use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedF
use crate::num::NonZero;
use crate::ops::Try;
use core::array;
-use core::mem::{ManuallyDrop, MaybeUninit};
+use core::mem::MaybeUninit;
use core::ops::ControlFlow;
/// 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")]
impl<I: fmt::Debug, P> fmt::Debug for Filter<I, P> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -64,52 +99,15 @@ where
fn next_chunk<const N: usize>(
&mut self,
) -> Result<[Self::Item; N], array::IntoIter<Self::Item, N>> {
- let mut array: [MaybeUninit<Self::Item>; N] = [const { MaybeUninit::uninit() }; N];
-
- struct Guard<'a, T> {
- array: &'a mut [MaybeUninit<T>],
- initialized: usize,
- }
-
- 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 fun = const {
+ if crate::mem::needs_drop::<I::Item>() {
+ array::iter_next_chunk::<I::Item, N>
+ } else {
+ Self::next_chunk_dropless::<N>
}
- }
-
- let mut guard = Guard { array: &mut array, initialized: 0 };
-
- 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) })
- }
- }
+ fun(self)
}
#[inline]
From 2be2d77c501730caedae8db2f0a0f02db57ae505 Mon Sep 17 00:00:00 2001
From: The 8472 <git@infinite-source.de>
Date: Mon, 24 Jun 2024 19:56:22 +0200
Subject: [PATCH 2/3] add comments explaining optimizations for
Filter::next_chunk
---
library/core/src/iter/adapters/filter.rs | 4 +++-
1 file changed, 3 insertions(+), 1 deletion(-)
diff --git a/library/core/src/iter/adapters/filter.rs b/library/core/src/iter/adapters/filter.rs
index 1c99f3938e270..ba49070329c22 100644
--- a/library/core/src/iter/adapters/filter.rs
+++ b/library/core/src/iter/adapters/filter.rs
@@ -41,8 +41,9 @@ where
let result = self.iter.try_for_each(|element| {
let idx = initialized;
+ // branchless index update combined with unconditionally copying the value even when
+ // it is filtered reduces branching and dependencies in the loop.
initialized = idx + (self.predicate)(&element) as usize;
-
// SAFETY: Loop conditions ensure the index is in bounds.
unsafe { array.get_unchecked_mut(idx) }.write(element);
@@ -99,6 +100,7 @@ where
fn next_chunk<const N: usize>(
&mut self,
) -> Result<[Self::Item; N], array::IntoIter<Self::Item, N>> {
+ // avoid codegen for the dead branch
let fun = const {
if crate::mem::needs_drop::<I::Item>() {
array::iter_next_chunk::<I::Item, N>
From 0d7aef9738b25ed4f84d11deb0add8d5345613b0 Mon Sep 17 00:00:00 2001
From: The 8472 <git@infinite-source.de>
Date: Tue, 25 Jun 2024 23:22:27 +0200
Subject: [PATCH 3/3] regression test for leaks in the the Filter::next_chunk
implementation
previously next_chunk would forget items rejected by the filter
---
library/core/tests/iter/adapters/filter.rs | 13 +++++++++++++
1 file changed, 13 insertions(+)
diff --git a/library/core/tests/iter/adapters/filter.rs b/library/core/tests/iter/adapters/filter.rs
index a2050d89d8564..167851e33336e 100644
--- a/library/core/tests/iter/adapters/filter.rs
+++ b/library/core/tests/iter/adapters/filter.rs
@@ -1,4 +1,5 @@
use core::iter::*;
+use std::rc::Rc;
#[test]
fn test_iterator_filter_count() {
@@ -50,3 +51,15 @@ fn test_double_ended_filter() {
assert_eq!(it.next().unwrap(), &2);
assert_eq!(it.next_back(), None);
}
+
+#[test]
+fn test_next_chunk_does_not_leak() {
+ let drop_witness: [_; 5] = std::array::from_fn(|_| Rc::new(()));
+
+ let v = (0..5).map(|i| drop_witness[i].clone()).collect::<Vec<_>>();
+ let _ = v.into_iter().filter(|_| false).next_chunk::<1>();
+
+ for ref w in drop_witness {
+ assert_eq!(Rc::strong_count(w), 1);
+ }
+}