spec next chunk for trustedlen

Author: bend-n

library/core/src/array/mod.rs

@@ -11,7 +11,7 @@ use crate::convert::Infallible;
 use crate::error::Error;
 use crate::hash::{self, Hash};
 use crate::intrinsics::transmute_unchecked;
-use crate::iter::{UncheckedIterator, repeat_n};
+use crate::iter::{TrustedLen, UncheckedIterator, repeat_n};
 use crate::marker::Destruct;
 use crate::mem::{self, ManuallyDrop, MaybeUninit};
 use crate::ops::{
@@ -964,6 +964,12 @@ impl<T: [const] Destruct> const Drop for Guard<'_, T> {
     }
 }
 
+pub(crate) fn iter_next_chunk<T, const N: usize>(
+    iter: &mut impl Iterator<Item = T>,
+) -> Result<[T; N], IntoIter<T, N>> {
+    iter.spec_next_chunk()
+}
+#[inline]
 /// Pulls `N` items from `iter` and returns them as an array. If the iterator
 /// yields fewer than `N` items, `Err` is returned containing an iterator over
 /// the already yielded items.
@@ -976,8 +982,7 @@ impl<T: [const] Destruct> const Drop for Guard<'_, T> {
 /// dropped.
 ///
 /// Used for [`Iterator::next_chunk`].
-#[inline]
-pub(crate) fn iter_next_chunk<T, const N: usize>(
+fn generic_iter_next_chunk<T, const N: usize>(
     iter: &mut impl Iterator<Item = T>,
 ) -> Result<[T; N], IntoIter<T, N>> {
     let mut array = [const { MaybeUninit::uninit() }; N];
@@ -994,6 +999,32 @@ pub(crate) fn iter_next_chunk<T, const N: usize>(
     }
 }
 
+pub(crate) trait SpecNextChunk<T, const N: usize>: Iterator<Item = T> {
+    fn spec_next_chunk(&mut self) -> Result<[T; N], IntoIter<T, N>>;
+}
+impl<I: Iterator<Item = T>, T, const N: usize> SpecNextChunk<T, N> for I {
+    default fn spec_next_chunk(&mut self) -> Result<[T; N], IntoIter<T, N>> {
+        generic_iter_next_chunk(self)
+    }
+}
+
+impl<I: Iterator<Item = T> + TrustedLen, T, const N: usize> SpecNextChunk<T, N> for I {
+    fn spec_next_chunk(&mut self) -> Result<[T; N], IntoIter<T, N>> {
+        if self.size_hint().0 < N {
+            let mut array = [const { MaybeUninit::uninit() }; N];
+            // SAFETY: Has to error out; trusted len means that there may
+            // only be less than N elements
+            let initialized =
+                unsafe { iter_next_chunk_erased(&mut array, self).unwrap_err_unchecked() };
+            // SAFETY: Only the first `initialized` elements were populated
+            return Err(unsafe { IntoIter::new_unchecked(array, 0..initialized) });
+        } else {
+            // SAFETY: must be at least N elements; safe to unwrap N elements.
+            return Ok(from_fn(|_| unsafe { self.next().unwrap_unchecked() }));
+        }
+    }
+}
+
 /// Version of [`iter_next_chunk`] using a passed-in slice in order to avoid
 /// needing to monomorphize for every array length.
 ///