diff --git a/library/core/benches/iter.rs b/library/core/benches/iter.rs index 9193c79bee8..60ef83223d1 100644 --- a/library/core/benches/iter.rs +++ b/library/core/benches/iter.rs @@ -404,7 +404,7 @@ fn bench_trusted_random_access_adapters(b: &mut Bencher) { /// Exercises the iter::Copied specialization for slice::Iter #[bench] -fn bench_copied_chunks(b: &mut Bencher) { +fn bench_next_chunk_copied(b: &mut Bencher) { let v = vec![1u8; 1024]; b.iter(|| { @@ -421,7 +421,7 @@ fn bench_copied_chunks(b: &mut Bencher) { /// Exercises the TrustedRandomAccess specialization in ArrayChunks #[bench] -fn bench_trusted_random_access_chunks(b: &mut Bencher) { +fn bench_next_chunk_trusted_random_access(b: &mut Bencher) { let v = vec![1u8; 1024]; b.iter(|| { @@ -437,3 +437,45 @@ fn bench_trusted_random_access_chunks(b: &mut Bencher) { .sum::>() }) } + +#[bench] +fn bench_next_chunk_filter_even(b: &mut Bencher) { + let a = (0..1024).next_chunk::<1024>().unwrap(); + + b.iter(|| black_box(&a).iter().filter(|&&i| i % 2 == 0).next_chunk::<32>()) +} + +#[bench] +fn bench_next_chunk_filter_predictably_true(b: &mut Bencher) { + let a = (0..1024).next_chunk::<1024>().unwrap(); + + b.iter(|| black_box(&a).iter().filter(|&&i| i < 100).next_chunk::<32>()) +} + +#[bench] +fn bench_next_chunk_filter_mostly_false(b: &mut Bencher) { + let a = (0..1024).next_chunk::<1024>().unwrap(); + + b.iter(|| black_box(&a).iter().filter(|&&i| i > 900).next_chunk::<32>()) +} + +#[bench] +fn bench_next_chunk_filter_map_even(b: &mut Bencher) { + let a = (0..1024).next_chunk::<1024>().unwrap(); + + b.iter(|| black_box(&a).iter().filter_map(|&i| (i % 2 == 0).then(|| i)).next_chunk::<32>()) +} + +#[bench] +fn bench_next_chunk_filter_map_predictably_true(b: &mut Bencher) { + let a = (0..1024).next_chunk::<1024>().unwrap(); + + b.iter(|| black_box(&a).iter().filter_map(|&i| (i < 100).then(|| i)).next_chunk::<32>()) +} + +#[bench] +fn bench_next_chunk_filter_map_mostly_false(b: &mut Bencher) { + let a = (0..1024).next_chunk::<1024>().unwrap(); + + b.iter(|| black_box(&a).iter().filter_map(|&i| (i > 900).then(|| i)).next_chunk::<32>()) +} diff --git a/library/core/src/iter/adapters/filter.rs b/library/core/src/iter/adapters/filter.rs index a0afaa326ad..723657b9e43 100644 --- a/library/core/src/iter/adapters/filter.rs +++ b/library/core/src/iter/adapters/filter.rs @@ -1,6 +1,9 @@ use crate::fmt; use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable}; use crate::ops::Try; +use core::array; +use core::mem::{ManuallyDrop, MaybeUninit}; +use core::ops::ControlFlow; /// An iterator that filters the elements of `iter` with `predicate`. /// @@ -56,6 +59,58 @@ fn next(&mut self) -> Option { self.iter.find(&mut self.predicate) } + #[inline] + fn next_chunk( + &mut self, + ) -> Result<[Self::Item; N], array::IntoIter> { + let mut array: [MaybeUninit; N] = MaybeUninit::uninit_array(); + + 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 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) }) + } + } + } + #[inline] fn size_hint(&self) -> (usize, Option) { let (_, upper) = self.iter.size_hint(); diff --git a/library/core/src/iter/adapters/filter_map.rs b/library/core/src/iter/adapters/filter_map.rs index 6bdf53f7fc9..693479977db 100644 --- a/library/core/src/iter/adapters/filter_map.rs +++ b/library/core/src/iter/adapters/filter_map.rs @@ -1,6 +1,7 @@ -use crate::fmt; use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable}; +use crate::mem::{ManuallyDrop, MaybeUninit}; use crate::ops::{ControlFlow, Try}; +use crate::{array, fmt}; /// An iterator that uses `f` to both filter and map elements from `iter`. /// @@ -61,6 +62,65 @@ fn next(&mut self) -> Option { self.iter.find_map(&mut self.f) } + #[inline] + fn next_chunk( + &mut self, + ) -> Result<[Self::Item; N], array::IntoIter> { + let mut array: [MaybeUninit; N] = MaybeUninit::uninit_array(); + + 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 mut guard = Guard { array: &mut array, initialized: 0 }; + + let result = self.iter.try_for_each(|element| { + let idx = guard.initialized; + let val = (self.f)(element); + guard.initialized = idx + val.is_some() as usize; + + // SAFETY: Loop conditions ensure the index is in bounds. + + unsafe { + let opt_payload_at = core::intrinsics::option_payload_ptr(&val); + let dst = guard.array.as_mut_ptr().add(idx); + crate::ptr::copy_nonoverlapping(opt_payload_at.cast(), dst, 1); + crate::mem::forget(val); + }; + + 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] fn size_hint(&self) -> (usize, Option) { let (_, upper) = self.iter.size_hint();