specialize slice_iter.copied().next_chunk()

1 files changed, 74 insertions, 0 deletions

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 @@ use crate::iter::adapters::{
     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 @@ where
         self.it.next().copied()
     }
 
+    fn next_chunk<const N: usize>(
+        &mut self,
+    ) -> Result<[Self::Item; N], array::IntoIter<Self::Item, N>>
+    where
+        Self: Sized,
+    {
+        <I as SpecNextChunk<'_, N, T>>::spec_next_chunk(&mut self.it)
+    }
+
     fn size_hint(&self) -> (usize, Option<usize>) {
         self.it.size_hint()
     }
@@ -166,3 +178,65 @@ where
     T: Copy,
 {
 }
+
+trait SpecNextChunk<'a, const N: usize, T: 'a>: Iterator<Item = &'a T>
+where
+    T: Copy,
+{
+    fn spec_next_chunk(&mut self) -> Result<[T; N], array::IntoIter<T, N>>;
+}
+
+impl<'a, const N: usize, I, T: 'a> SpecNextChunk<'a, N, T> for I
+where
+    I: Iterator<Item = &'a T>,
+    T: Copy,
+{
+    default fn spec_next_chunk(&mut self) -> Result<[T; N], array::IntoIter<T, N>> {
+        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<T, N>> {
+        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))
+        }
+    }
+}