diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index 6cca2e6358b63..f71414a5a7eaa 100644 --- a/library/core/src/array/mod.rs +++ b/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::{ @@ -979,6 +979,13 @@ impl const Drop for Guard<'_, T> { #[inline] pub(crate) fn iter_next_chunk( iter: &mut impl Iterator, +) -> Result<[T; N], IntoIter> { + iter.spec_next_chunk() +} + +#[inline] +fn generic_iter_next_chunk( + iter: &mut impl Iterator, ) -> Result<[T; N], IntoIter> { let mut array = [const { MaybeUninit::uninit() }; N]; let r = iter_next_chunk_erased(&mut array, iter); @@ -994,6 +1001,33 @@ pub(crate) fn iter_next_chunk( } } +pub(crate) trait SpecNextChunk: Iterator { + fn spec_next_chunk(&mut self) -> Result<[T; N], IntoIter>; +} +impl, T, const N: usize> SpecNextChunk for I { + #[inline] + default fn spec_next_chunk(&mut self) -> Result<[T; N], IntoIter> { + generic_iter_next_chunk(self) + } +} + +impl + TrustedLen, T, const N: usize> SpecNextChunk for I { + fn spec_next_chunk(&mut self) -> Result<[T; N], IntoIter> { + if self.size_hint().0 < N { + let mut array = [const { MaybeUninit::uninit() }; N]; + let initialized = + // SAFETY: Has to error out; trusted len means + // that there may only be less than N elements + unsafe { iter_next_chunk_erased(&mut array, self).unwrap_err_unchecked() }; + // SAFETY: Only the first `initialized` elements were populated + Err(unsafe { IntoIter::new_unchecked(array, 0..initialized) }) + } else { + // SAFETY: must be at least N elements; safe to unwrap N elements. + 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. ///