Add Iterator::array_chunks()

library/core/src/array/iter.rs

@@ -84,6 +84,16 @@ impl<T, const N: usize> IntoIter<T, N> {
         IntoIterator::into_iter(array)
     }
 
+    /// Creates a new iterator from a partially initalized array.
+    ///
+    /// # Safety
+    ///
+    /// The caller must guarantee that all and only the `alive` elements of
+    /// `data` are initialized.
+    pub(crate) unsafe fn with_partial(data: [MaybeUninit<T>; N], alive: Range<usize>) -> Self {
+        Self { data, alive }
+    }
+
     /// Creates an iterator over the elements in a partially-initialized buffer.
     ///
     /// If you have a fully-initialized array, then use [`IntoIterator`].

library/core/src/iter/adapters/array_chunks.rs

@@ -0,0 +1,362 @@
+use crate::array;
+use crate::iter::{Fuse, FusedIterator, Iterator, TrustedLen};
+use crate::mem;
+use crate::mem::MaybeUninit;
+use crate::ops::{ControlFlow, Try};
+use crate::ptr;
+
+/// An iterator over `N` elements of the iterator at a time.
+///
+/// The chunks do not overlap. If `N` does not divide the length of the
+/// iterator, then the last up to `N-1` elements will be omitted.
+///
+/// This `struct` is created by the [`array_chunks`][Iterator::array_chunks]
+/// method on [`Iterator`]. See its documentation for more.
+#[derive(Debug, Clone)]
+#[must_use = "iterators are lazy and do nothing unless consumed"]
+#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")]
+pub struct ArrayChunks<I: Iterator, const N: usize> {
+    iter: Fuse<I>,
+    remainder: Option<array::IntoIter<I::Item, N>>,
+}
+
+impl<I, const N: usize> ArrayChunks<I, N>
+where
+    I: Iterator,
+{
+    pub(in crate::iter) fn new(iter: I) -> Self {
+        assert!(N != 0, "chunk size must be non-zero");
+        Self { iter: iter.fuse(), remainder: None }
+    }
+
+    /// Returns an iterator over the remaining elements of the original iterator
+    /// that are not going to be returned by this iterator. The returned
+    /// iterator will yield at most `N-1` elements.
+    #[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")]
+    #[inline]
+    pub fn into_remainder(self) -> Option<array::IntoIter<I::Item, N>> {
+        self.remainder
+    }
+}
+
+#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")]
+impl<I, const N: usize> Iterator for ArrayChunks<I, N>
+where
+    I: Iterator,
+{
+    type Item = [I::Item; N];
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        let mut array = MaybeUninit::uninit_array();
+        // SAFETY: `array` will still be valid if `guard` is dropped.
+        let mut guard = unsafe { FrontGuard::new(&mut array) };
+
+        for slot in array.iter_mut() {
+            match self.iter.next() {
+                Some(item) => {
+                    slot.write(item);
+                    guard.init += 1;
+                }
+                None => {
+                    if guard.init > 0 {
+                        let init = guard.init;
+                        mem::forget(guard);
+                        self.remainder = {
+                            // SAFETY: `array` was initialized with `init` elements.
+                            Some(unsafe { array::IntoIter::with_partial(array, 0..init) })
+                        };
+                    }
+                    return None;
+                }
+            }
+        }
+
+        mem::forget(guard);
+        // SAFETY: All elements of the array were populated in the loop above.
+        Some(unsafe { MaybeUninit::array_assume_init(array) })
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let (lower, upper) = self.iter.size_hint();
+        // Keep infinite iterator size hint lower bound as `usize::MAX`. This
+        // is required to implement `TrustedLen`.
+        if lower == usize::MAX {
+            return (lower, upper);
+        }
+        (lower / N, upper.map(|n| n / N))
+    }
+
+    #[inline]
+    fn count(self) -> usize {
+        self.iter.count() / N
+    }
+
+    fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Output = B>,
+    {
+        let mut array = MaybeUninit::uninit_array();
+        // SAFETY: `array` will still be valid if `guard` is dropped.
+        let mut guard = unsafe { FrontGuard::new(&mut array) };
+
+        let result = self.iter.try_fold(init, |mut acc, item| {
+            // SAFETY: `init` starts at 0, increases by one each iteration and
+            // is reset to 0 once it reaches N.
+            unsafe { array.get_unchecked_mut(guard.init) }.write(item);
+            guard.init += 1;
+            if guard.init == N {
+                guard.init = 0;
+                let array = mem::replace(&mut array, MaybeUninit::uninit_array());
+                // SAFETY: the condition above asserts that all elements are
+                // initialized.
+                let item = unsafe { MaybeUninit::array_assume_init(array) };
+                acc = f(acc, item)?;
+            }
+            R::from_output(acc)
+        });
+        match result.branch() {
+            ControlFlow::Continue(o) => {
+                if guard.init > 0 {
+                    let init = guard.init;
+                    mem::forget(guard);
+                    // SAFETY: `array` was initialized with `init` elements.
+                    self.remainder = Some(unsafe { array::IntoIter::with_partial(array, 0..init) });
+                }
+                R::from_output(o)
+            }
+            ControlFlow::Break(r) => R::from_residual(r),
+        }
+    }
+
+    fn fold<B, F>(self, init: B, mut f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        let mut array = MaybeUninit::uninit_array();
+        // SAFETY: `array` will still be valid if `guard` is dropped.
+        let mut guard = unsafe { FrontGuard::new(&mut array) };
+
+        self.iter.fold(init, |mut acc, item| {
+            // SAFETY: `init` starts at 0, increases by one each iteration and
+            // is reset to 0 once it reaches N.
+            unsafe { array.get_unchecked_mut(guard.init) }.write(item);
+            guard.init += 1;
+            if guard.init == N {
+                guard.init = 0;
+                let array = mem::replace(&mut array, MaybeUninit::uninit_array());
+                // SAFETY: the condition above asserts that all elements are
+                // initialized.
+                let item = unsafe { MaybeUninit::array_assume_init(array) };
+                acc = f(acc, item);
+            }
+            acc
+        })
+    }
+}
+
+/// A guard for an array where elements are filled from the left.
+struct FrontGuard<T, const N: usize> {
+    /// A pointer to the array that is being filled. We need to use a raw
+    /// pointer here because of the lifetime issues in the fold implementations.
+    ptr: *mut T,
+    /// The number of *initialized* elements.
+    init: usize,
+}
+
+impl<T, const N: usize> FrontGuard<T, N> {
+    unsafe fn new(array: &mut [MaybeUninit<T>; N]) -> Self {
+        Self { ptr: MaybeUninit::slice_as_mut_ptr(array), init: 0 }
+    }
+}
+
+impl<T, const N: usize> Drop for FrontGuard<T, N> {
+    fn drop(&mut self) {
+        debug_assert!(self.init <= N);
+        // SAFETY: This raw slice will only contain the initialized objects
+        // within the buffer.
+        unsafe {
+            let slice = ptr::slice_from_raw_parts_mut(self.ptr, self.init);
+            ptr::drop_in_place(slice);
+        }
+    }
+}
+
+#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")]
+impl<I, const N: usize> DoubleEndedIterator for ArrayChunks<I, N>
+where
+    I: DoubleEndedIterator + ExactSizeIterator,
+{
+    #[inline]
+    fn next_back(&mut self) -> Option<Self::Item> {
+        // We are iterating from the back we need to first handle the remainder.
+        self.next_back_remainder()?;
+
+        let mut array = MaybeUninit::uninit_array();
+        // SAFETY: `array` will still be valid if `guard` is dropped.
+        let mut guard = unsafe { BackGuard::new(&mut array) };
+
+        for slot in array.iter_mut().rev() {
+            slot.write(self.iter.next_back()?);
+            guard.uninit -= 1;
+        }
+
+        mem::forget(guard);
+        // SAFETY: All elements of the array were populated in the loop above.
+        Some(unsafe { MaybeUninit::array_assume_init(array) })
+    }
+
+    fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Output = B>,
+    {
+        // We are iterating from the back we need to first handle the remainder.
+        if self.next_back_remainder().is_none() {
+            return R::from_output(init);
+        }
+
+        let mut array = MaybeUninit::uninit_array();
+        // SAFETY: `array` will still be valid if `guard` is dropped.
+        let mut guard = unsafe { BackGuard::new(&mut array) };
+
+        self.iter.try_rfold(init, |mut acc, item| {
+            guard.uninit -= 1;
+            // SAFETY: `uninit` starts at N, decreases by one each iteration and
+            // is reset to N once it reaches 0.
+            unsafe { array.get_unchecked_mut(guard.uninit) }.write(item);
+            if guard.uninit == 0 {
+                guard.uninit = N;
+                let array = mem::replace(&mut array, MaybeUninit::uninit_array());
+                // SAFETY: the condition above asserts that all elements are
+                // initialized.
+                let item = unsafe { MaybeUninit::array_assume_init(array) };
+                acc = f(acc, item)?;
+            }
+            R::from_output(acc)
+        })
+    }
+
+    fn rfold<B, F>(mut self, init: B, mut f: F) -> B
+    where
+        Self: Sized,
+        F: FnMut(B, Self::Item) -> B,
+    {
+        // We are iterating from the back we need to first handle the remainder.
+        if self.next_back_remainder().is_none() {
+            return init;
+        }
+
+        let mut array = MaybeUninit::uninit_array();
+
+        // SAFETY: `array` will still be valid if `guard` is dropped.
+        let mut guard = unsafe { BackGuard::new(&mut array) };
+
+        self.iter.rfold(init, |mut acc, item| {
+            guard.uninit -= 1;
+            // SAFETY: `uninit` starts at N, decreases by one each iteration and
+            // is reset to N once it reaches 0.
+            unsafe { array.get_unchecked_mut(guard.uninit) }.write(item);
+            if guard.uninit == 0 {
+                guard.uninit = N;
+                let array = mem::replace(&mut array, MaybeUninit::uninit_array());
+                // SAFETY: the condition above asserts that all elements are
+                // initialized.
+                let item = unsafe { MaybeUninit::array_assume_init(array) };
+                acc = f(acc, item);
+            }
+            acc
+        })
+    }
+}
+
+impl<I, const N: usize> ArrayChunks<I, N>
+where
+    I: DoubleEndedIterator + ExactSizeIterator,
+{
+    #[inline]
+    fn next_back_remainder(&mut self) -> Option<()> {
+        // We use the `ExactSizeIterator` implementation of the underlying
+        // iterator to know how many remaining elements there are.
+        let rem = self.iter.len() % N;
+        if rem == 0 {
+            return Some(());
+        }
+
+        let mut array = MaybeUninit::uninit_array();
+
+        // SAFETY: The array will still be valid if `guard` is dropped and
+        // it is forgotten otherwise.
+        let mut guard = unsafe { FrontGuard::new(&mut array) };
+
+        // SAFETY: `rem` is in the range 1..N based on how it is calculated.
+        for slot in unsafe { array.get_unchecked_mut(..rem) }.iter_mut() {
+            slot.write(self.iter.next_back()?);
+            guard.init += 1;
+        }
+
+        let init = guard.init;
+        mem::forget(guard);
+        // SAFETY: `array` was initialized with exactly `init` elements.
+        self.remainder = unsafe {
+            array.get_unchecked_mut(..init).reverse();
+            Some(array::IntoIter::with_partial(array, 0..init))
+        };
+        Some(())
+    }
+}
+
+/// A guard for an array where elements are filled from the right.
+struct BackGuard<T, const N: usize> {
+    /// A pointer to the array that is being filled. We need to use a raw
+    /// pointer here because of the lifetime issues in the rfold implementations.
+    ptr: *mut T,
+    /// The number of *uninitialized* elements.
+    uninit: usize,
+}
+
+impl<T, const N: usize> BackGuard<T, N> {
+    unsafe fn new(array: &mut [MaybeUninit<T>; N]) -> Self {
+        Self { ptr: MaybeUninit::slice_as_mut_ptr(array), uninit: N }
+    }
+}
+
+impl<T, const N: usize> Drop for BackGuard<T, N> {
+    fn drop(&mut self) {
+        debug_assert!(self.uninit <= N);
+        // SAFETY: This raw slice will only contain the initialized objects
+        // within the buffer.
+        unsafe {
+            let ptr = self.ptr.offset(self.uninit as isize);
+            let slice = ptr::slice_from_raw_parts_mut(ptr, N - self.uninit);
+            ptr::drop_in_place(slice);
+        }
+    }
+}
+
+#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")]
+impl<I, const N: usize> FusedIterator for ArrayChunks<I, N> where I: FusedIterator {}
+
+#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")]
+impl<I, const N: usize> ExactSizeIterator for ArrayChunks<I, N>
+where
+    I: ExactSizeIterator,
+{
+    #[inline]
+    fn len(&self) -> usize {
+        self.iter.len() / N
+    }
+
+    #[inline]
+    fn is_empty(&self) -> bool {
+        self.iter.len() / N == 0
+    }
+}
+
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<I, const N: usize> TrustedLen for ArrayChunks<I, N> where I: TrustedLen {}

library/core/src/iter/adapters/mod.rs

@@ -1,6 +1,7 @@
 use crate::iter::{InPlaceIterable, Iterator};
 use crate::ops::{ControlFlow, Try};
 
+mod array_chunks;
 mod chain;
 mod cloned;
 mod copied;
@@ -31,6 +32,9 @@ pub use self::{
     scan::Scan, skip::Skip, skip_while::SkipWhile, take::Take, take_while::TakeWhile, zip::Zip,
 };
 
+#[unstable(feature = "iter_array_chunks", reason = "recently added", issue = "none")]
+pub use self::array_chunks::ArrayChunks;
+
 #[stable(feature = "iter_cloned", since = "1.1.0")]
 pub use self::cloned::Cloned;