Add Iterator::array_chunks method

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

@@ -0,0 +1,73 @@
+use core::iter::FusedIterator;
+use core::mem::{self, MaybeUninit};
+use core::ptr;
+
+/// An iterator that yields the elements of another iterator in
+/// chunks of size `N`.
+///
+/// This `struct` is created by the [`array_chunks`] method on [`Iterator`]. See
+/// its documentation for more.
+///
+/// [`array_chunks`]: Iterator::array_chunks
+#[unstable(feature = "iter_array_chunks", issue = "none")]
+#[derive(Debug)]
+pub struct ArrayChunks<I: Iterator, const N: usize> {
+    iter: I,
+    buffer: [MaybeUninit<I::Item>; N],
+    init: usize,
+}
+
+#[unstable(feature = "iter_array_chunks", issue = "none")]
+impl<I: Iterator, const N: usize> Drop for ArrayChunks<I, N> {
+    fn drop(&mut self) {
+        // SAFETY: This is safe: `remainder_mut` returns exactly the sub-slice
+        // of elements that were initialized but not yielded and so have yet
+        // to be dropped.
+        unsafe {
+            ptr::drop_in_place(self.remainder_mut());
+        }
+    }
+}
+
+impl<I: Iterator, const N: usize> ArrayChunks<I, N> {
+    pub(in crate::iter) fn new(iter: I) -> Self {
+        Self { iter, init: 0, buffer: MaybeUninit::uninit_array() }
+    }
+
+    /// Returns the remainder of the elements yielded by the original
+    /// iterator that were insufficient to fill another chunk. The
+    /// returned slice has at most `N-1` elements.
+    #[unstable(feature = "iter_array_chunks", issue = "none")]
+    pub fn remainder(&self) -> &[I::Item] {
+        // SAFETY: We know that all elements before `init` are properly initialized.
+        unsafe { MaybeUninit::slice_assume_init_ref(&self.buffer[..self.init]) }
+    }
+
+    /// Returns the remainder of the elements yielded by the original
+    /// iterator that were insufficient to fill another chunk. The
+    /// returned slice has at most `N-1` elements.
+    #[unstable(feature = "iter_array_chunks", issue = "none")]
+    pub fn remainder_mut(&mut self) -> &mut [I::Item] {
+        // SAFETY: We know that all elements before `init` are properly initialized.
+        unsafe { MaybeUninit::slice_assume_init_mut(&mut self.buffer[..self.init]) }
+    }
+}
+
+#[unstable(feature = "iter_array_chunks", issue = "none")]
+impl<I: Iterator, const N: usize> Iterator for ArrayChunks<I, N> {
+    type Item = [I::Item; N];
+
+    fn next(&mut self) -> Option<Self::Item> {
+        while self.init < N {
+            self.buffer[self.init] = MaybeUninit::new(self.iter.next()?);
+            self.init += 1;
+        }
+        self.init = 0;
+        // SAFETY: This is safe: `MaybeUninit<T>` is guaranteed to have the same layout
+        // as `T` and the entire array has just been initialized.
+        unsafe { Some(mem::transmute_copy(&self.buffer)) }
+    }
+}
+
+#[unstable(feature = "iter_array_chunks", issue = "none")]
+impl<I: FusedIterator, const N: usize> FusedIterator for ArrayChunks<I, N> {}

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;
@@ -30,6 +31,9 @@ pub use self::{
     scan::Scan, skip::Skip, skip_while::SkipWhile, take::Take, take_while::TakeWhile, zip::Zip,
 };
 
+#[unstable(feature = "iter_array_chunks", issue = "none")]
+pub use self::array_chunks::ArrayChunks;
+
 #[stable(feature = "iter_cloned", since = "1.1.0")]
 pub use self::cloned::Cloned;
 

library/core/src/iter/mod.rs

@@ -393,6 +393,8 @@ pub use self::traits::{
 
 #[unstable(feature = "iter_zip", issue = "83574")]
 pub use self::adapters::zip;
+#[unstable(feature = "iter_array_chunks", issue = "none")]
+pub use self::adapters::ArrayChunks;
 #[stable(feature = "iter_cloned", since = "1.1.0")]
 pub use self::adapters::Cloned;
 #[stable(feature = "iter_copied", since = "1.36.0")]

library/core/src/iter/traits/iterator.rs

@@ -5,6 +5,7 @@
 use crate::cmp::{self, Ordering};
 use crate::ops::{ControlFlow, Try};
 
+use super::super::ArrayChunks;
 use super::super::TrustedRandomAccessNoCoerce;
 use super::super::{Chain, Cloned, Copied, Cycle, Enumerate, Filter, FilterMap, Fuse};
 use super::super::{FlatMap, Flatten};
@@ -3468,6 +3469,36 @@ pub trait Iterator {
     {
         unreachable!("Always specialized");
     }
+
+    /// Creates an iterator which yields arrays of `N` elements yielded by
+    /// the original iterator.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `N` is zero.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(iter_array_chunks)]
+    /// let mut iter = (0..10).array_chunks::<3>();
+    ///
+    /// assert_eq!(iter.next(), Some([0, 1, 2]));
+    /// assert_eq!(iter.next(), Some([3, 4, 5]));
+    /// assert_eq!(iter.next(), Some([6, 7, 8]));
+    /// assert_eq!(iter.next(), None);
+    ///
+    /// assert_eq!(iter.remainder(), &[9]);
+    /// ```
+    #[inline]
+    #[unstable(feature = "iter_array_chunks", issue = "none")]
+    fn array_chunks<const N: usize>(self) -> ArrayChunks<Self, N>
+    where
+        Self: Sized,
+    {
+        assert_ne!(N, 0);
+        ArrayChunks::new(self)
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]

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

@@ -0,0 +1,66 @@
+use core::cell::Cell;
+use core::iter::*;
+
+#[derive(Debug)]
+struct DropBomb<'a> {
+    dropped: bool,
+    counter: &'a Cell<usize>,
+}
+
+impl Drop for DropBomb<'_> {
+    fn drop(&mut self) {
+        if self.dropped {
+            panic!("double dropped!");
+        }
+        self.dropped = true;
+        self.counter.set(self.counter.get() + 1);
+    }
+}
+
+#[test]
+fn test_iterator_array_chunks_remainder() {
+    let mut iter = (0..=10).array_chunks::<4>();
+    assert_eq!(iter.remainder(), &[]);
+    assert_eq!(iter.remainder_mut(), &[]);
+    assert_eq!(iter.next(), Some([0, 1, 2, 3]));
+    assert_eq!(iter.remainder(), &[]);
+    assert_eq!(iter.remainder_mut(), &[]);
+    assert_eq!(iter.next(), Some([4, 5, 6, 7]));
+    assert_eq!(iter.remainder(), &[]);
+    assert_eq!(iter.remainder_mut(), &[]);
+    assert_eq!(iter.next(), None);
+    assert_eq!(iter.remainder(), &[8, 9, 10]);
+    assert_eq!(iter.remainder_mut(), &[8, 9, 10]);
+}
+
+#[test]
+fn test_iterator_array_chunks_drop() {
+    let counter = Cell::new(0);
+    let create =
+        |n| (0..n).map(|_| DropBomb { dropped: false, counter: &counter }).array_chunks::<3>();
+
+    let iter = create(5);
+    assert_eq!(counter.get(), 0);
+    drop(iter);
+    assert_eq!(counter.get(), 0);
+
+    let mut iter = create(3);
+    counter.set(0);
+    iter.next();
+    assert_eq!(counter.get(), 3);
+    assert!(iter.next().is_none());
+    assert_eq!(counter.get(), 3);
+    assert_eq!(iter.remainder().len(), 0);
+    drop(iter);
+    assert_eq!(counter.get(), 3);
+
+    let mut iter = create(5);
+    counter.set(0);
+    iter.next();
+    assert_eq!(counter.get(), 3);
+    assert!(iter.next().is_none());
+    assert_eq!(counter.get(), 3);
+    assert_eq!(iter.remainder().len(), 2);
+    drop(iter);
+    assert_eq!(counter.get(), 5);
+}

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

@@ -1,3 +1,4 @@
+mod array_chunks;
 mod chain;
 mod cloned;
 mod copied;

library/core/tests/lib.rs

@@ -45,6 +45,7 @@
 #![feature(slice_partition_dedup)]
 #![feature(int_log)]
 #![feature(iter_advance_by)]
+#![feature(iter_array_chunks)]
 #![feature(iter_partition_in_place)]
 #![feature(iter_intersperse)]
 #![feature(iter_is_partitioned)]