Add Iterator::array_chunks method

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

@@ -0,0 +1,146 @@
+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> {}
+
+/// An iterator that yields the elements of another iterator in
+/// chunks of size `N` starting from the end.
+///
+/// This `struct` is created by the [`array_rchunks`] method on [`Iterator`]. See
+/// its documentation for more.
+///
+/// [`array_rchunks`]: Iterator::array_rchunks
+#[unstable(feature = "iter_array_chunks", issue = "none")]
+#[derive(Debug)]
+pub struct ArrayRChunks<I: DoubleEndedIterator, const N: usize> {
+ iter: I,
+ buffer: [MaybeUninit<I::Item>; N],
+ init: usize,
+}
+
+#[unstable(feature = "iter_array_chunks", issue = "none")]
+impl<I: DoubleEndedIterator, const N: usize> Drop for ArrayRChunks<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: DoubleEndedIterator, const N: usize> ArrayRChunks<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 after `init` are properly initialized.
+ unsafe { MaybeUninit::slice_assume_init_ref(&self.buffer[(N - 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 after `init` are properly initialized.
+ unsafe { MaybeUninit::slice_assume_init_mut(&mut self.buffer[(N - self.init)..]) }
+ }
+}
+
+#[unstable(feature = "iter_array_chunks", issue = "none")]
+impl<I: DoubleEndedIterator, const N: usize> Iterator for ArrayRChunks<I, N> {
+ type Item = [I::Item; N];
+
+ fn next(&mut self) -> Option<Self::Item> {
+ while self.init < N {
+ self.buffer[N - self.init - 1] = MaybeUninit::new(self.iter.next_back()?);
+ 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, const N: usize> FusedIterator for ArrayRChunks<I, N> where
+ I: DoubleEndedIterator + FusedIterator
+{
+}

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, ArrayRChunks};
+
 #[stable(feature = "iter_cloned", since = "1.1.0")]
 pub use self::cloned::Cloned;
 

library/core/src/iter/mod.rs

@@ -409,6 +409,8 @@ pub use self::adapters::StepBy;
 pub use self::adapters::TrustedRandomAccess;
 #[unstable(feature = "trusted_random_access", issue = "none")]
 pub use self::adapters::TrustedRandomAccessNoCoerce;
+#[unstable(feature = "iter_array_chunks", issue = "none")]
+pub use self::adapters::{ArrayChunks, ArrayRChunks};
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::adapters::{
  Chain, Cycle, Enumerate, Filter, FilterMap, FlatMap, Fuse, Inspect, Map, Peekable, Rev, Scan,

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

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

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

@@ -0,0 +1,114 @@
+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);
+}
+
+#[test]
+fn test_iterator_array_rchunks_remainder() {
+ let mut iter = (0..=10).array_rchunks::<4>();
+ assert_eq!(iter.remainder(), &[]);
+ assert_eq!(iter.remainder_mut(), &[]);
+ assert_eq!(iter.next(), Some([7, 8, 9, 10]));
+ assert_eq!(iter.remainder(), &[]);
+ assert_eq!(iter.remainder_mut(), &[]);
+ assert_eq!(iter.next(), Some([3, 4, 5, 6]));
+ assert_eq!(iter.remainder(), &[]);
+ assert_eq!(iter.remainder_mut(), &[]);
+ assert_eq!(iter.next(), None);
+ assert_eq!(iter.remainder(), &[0, 1, 2]);
+ assert_eq!(iter.remainder_mut(), &[0, 1, 2]);
+}
+
+#[test]
+fn test_iterator_array_rchunks_drop() {
+ let counter = Cell::new(0);
+ let create =
+ |n| (0..n).map(|_| DropBomb { dropped: false, counter: &counter }).array_rchunks::<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)]