From 873a18e2213a506d4cf6357e5faef0990ec08683 Mon Sep 17 00:00:00 2001 From: The 8472 Date: Mon, 17 Oct 2022 22:49:32 +0200 Subject: [PATCH] specialize slice_iter.copied().next_chunk() --- library/core/src/iter/adapters/copied.rs | 74 ++++++++++++++++++++++++ 1 file changed, 74 insertions(+) diff --git a/library/core/src/iter/adapters/copied.rs b/library/core/src/iter/adapters/copied.rs index f9bfd77d7fb..62d3afb8160 100644 --- a/library/core/src/iter/adapters/copied.rs +++ b/library/core/src/iter/adapters/copied.rs @@ -2,7 +2,10 @@ zip::try_get_unchecked, TrustedRandomAccess, TrustedRandomAccessNoCoerce, }; use crate::iter::{FusedIterator, TrustedLen}; +use crate::mem::MaybeUninit; +use crate::mem::SizedTypeProperties; use crate::ops::Try; +use crate::{array, ptr}; /// An iterator that copies the elements of an underlying iterator. /// @@ -44,6 +47,15 @@ fn next(&mut self) -> Option { self.it.next().copied() } + fn next_chunk( + &mut self, + ) -> Result<[Self::Item; N], array::IntoIter> + where + Self: Sized, + { + >::spec_next_chunk(&mut self.it) + } + fn size_hint(&self) -> (usize, Option) { self.it.size_hint() } @@ -166,3 +178,65 @@ unsafe impl<'a, I, T: 'a> TrustedLen for Copied T: Copy, { } + +trait SpecNextChunk<'a, const N: usize, T: 'a>: Iterator +where + T: Copy, +{ + fn spec_next_chunk(&mut self) -> Result<[T; N], array::IntoIter>; +} + +impl<'a, const N: usize, I, T: 'a> SpecNextChunk<'a, N, T> for I +where + I: Iterator, + T: Copy, +{ + default fn spec_next_chunk(&mut self) -> Result<[T; N], array::IntoIter> { + array::iter_next_chunk(&mut self.map(|e| *e)) + } +} + +impl<'a, const N: usize, T: 'a> SpecNextChunk<'a, N, T> for crate::slice::Iter<'a, T> +where + T: Copy, +{ + fn spec_next_chunk(&mut self) -> Result<[T; N], array::IntoIter> { + let mut raw_array = MaybeUninit::uninit_array(); + + let len = self.len(); + + if T::IS_ZST { + if len < N { + let _ = self.advance_by(len); + // SAFETY: ZSTs can be conjured ex nihilo; only the amount has to be correct + return Err(unsafe { array::IntoIter::new_unchecked(raw_array, 0..len) }); + } + + let _ = self.advance_by(N); + // SAFETY: ditto + return Ok(unsafe { MaybeUninit::array_assume_init(raw_array) }); + } + + if len < N { + // SAFETY: `len` indicates that this many elements are available and we just checked that + // it fits into the array. + unsafe { + ptr::copy_nonoverlapping( + self.as_ref().as_ptr(), + raw_array.as_mut_ptr() as *mut T, + len, + ); + let _ = self.advance_by(len); + return Err(array::IntoIter::new_unchecked(raw_array, 0..len)); + } + } + + // SAFETY: `len` is larger than the array size. Copy a fixed amount here to fully initialize + // the array. + unsafe { + ptr::copy_nonoverlapping(self.as_ref().as_ptr(), raw_array.as_mut_ptr() as *mut T, N); + let _ = self.advance_by(N); + Ok(MaybeUninit::array_assume_init(raw_array)) + } + } +}