From 133e7b10a45bddd4ef5ba265d848ce0a1976e8ae Mon Sep 17 00:00:00 2001 From: The 8472 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 ca23d1b13a88..1c99f3938e27 100644 --- a/library/core/src/iter/adapters/filter.rs +++ b/library/core/src/iter/adapters/filter.rs @@ -3,7 +3,7 @@ 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 @@ pub(in crate::iter) fn new(iter: I, predicate: P) -> Filter { } } +impl Filter +where + I: Iterator, + P: FnMut(&I::Item) -> bool, +{ + #[inline] + fn next_chunk_dropless( + &mut self, + ) -> Result<[I::Item; N], array::IntoIter> { + let mut array: [MaybeUninit; 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 fmt::Debug for Filter { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { @@ -64,52 +99,15 @@ fn next(&mut self) -> Option { fn next_chunk( &mut self, ) -> Result<[Self::Item; N], array::IntoIter> { - let mut array: [MaybeUninit; N] = [const { MaybeUninit::uninit() }; N]; - - struct Guard<'a, T> { - array: &'a mut [MaybeUninit], - initialized: usize, - } - - impl Drop for Guard<'_, T> { - #[inline] - fn drop(&mut self) { - if const { crate::mem::needs_drop::() } { - // 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::() { + array::iter_next_chunk:: + } else { + Self::next_chunk_dropless:: } - } + }; - 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 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 1c99f3938e27..ba49070329c2 100644 --- a/library/core/src/iter/adapters/filter.rs +++ b/library/core/src/iter/adapters/filter.rs @@ -41,8 +41,9 @@ fn next_chunk_dropless( 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 @@ fn next(&mut self) -> Option { fn next_chunk( &mut self, ) -> Result<[Self::Item; N], array::IntoIter> { + // avoid codegen for the dead branch let fun = const { if crate::mem::needs_drop::() { array::iter_next_chunk:: From 0d7aef9738b25ed4f84d11deb0add8d5345613b0 Mon Sep 17 00:00:00 2001 From: The 8472 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 a2050d89d856..167851e33336 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::>(); + let _ = v.into_iter().filter(|_| false).next_chunk::<1>(); + + for ref w in drop_witness { + assert_eq!(Rc::strong_count(w), 1); + } +}