optimize slice::Iter::next_chunk

Author: bend-n

library/core/src/slice/iter.rs

@@ -139,7 +139,7 @@ impl<'a, T> Iter<'a, T> {
     }
 }
 
-iterator! {struct Iter -> *const T, &'a T, const, {/* no mut */}, as_ref, {
+iterator! {struct Iter -> *const T, &'a T, const, {/* no mut */}, as_ref, each_ref, {
     fn is_sorted_by<F>(self, mut compare: F) -> bool
     where
         Self: Sized,
@@ -368,7 +368,7 @@ impl<T> AsRef<[T]> for IterMut<'_, T> {
 //     }
 // }
 
-iterator! {struct IterMut -> *mut T, &'a mut T, mut, {mut}, as_mut, {}}
+iterator! {struct IterMut -> *mut T, &'a mut T, mut, {mut}, as_mut, each_mut, {}}
 
 /// An internal abstraction over the splitting iterators, so that
 /// splitn, splitn_mut etc can be implemented once.

library/core/src/slice/iter/macros.rs

@@ -68,6 +68,7 @@ macro_rules! iterator {
         $raw_mut:tt,
         {$( $mut_:tt )?},
         $into_ref:ident,
+        $array_ref:ident,
         {$($extra:tt)*}
     ) => {
         impl<'a, T> $name<'a, T> {
@@ -190,6 +191,29 @@ macro_rules! iterator {
                 }
             }
 
+            fn next_chunk<const N:usize>(&mut self) -> Result<[$elem; N], crate::array::IntoIter<$elem, N>> {
+                if T::IS_ZST {
+                    return crate::array::iter_next_chunk(self);
+                }
+                let len = len!(self);
+                if len >= N {
+                    // SAFETY: we are just getting an array of [T; N] and moving the pointer over a little
+                    let r = unsafe { self.post_inc_start(N).cast_array().$into_ref() }
+                        .$array_ref(); // must convert &[T; N] to [&T; N]
+                    Ok(r)
+                } else {
+                    // cant use $array_ref because theres no builtin for &mut [MU<T>; N] -> [&mut MU<T>; N]
+                    // cant use copy_nonoverlapping as the $elem is of type &{mut} T instead of T
+                    let mut a = [const { crate::mem::MaybeUninit::<$elem>::uninit() }; N];
+                    for into in (&mut a).into_iter().take(len) {
+                        // SAFETY: take(n) limits to remainder (slice produces worse codegen)
+                        into.write(unsafe { self.post_inc_start(1).$into_ref() });
+                    }
+                    // SAFETY: we just initialized elements 0..len
+                    unsafe { Err(crate::array::IntoIter::new_unchecked(a, 0..len)) }
+                }
+            }
+
             #[inline]
             fn size_hint(&self) -> (usize, Option<usize>) {
                 let exact = len!(self);