Added next_array and collect_array.

Author: orlp

src/lib.rs

@@ -206,6 +206,7 @@ mod merge_join;
 mod minmax;
 #[cfg(feature = "use_alloc")]
 mod multipeek_impl;
+mod next_array;
 mod pad_tail;
 #[cfg(feature = "use_alloc")]
 mod peek_nth;
@@ -1914,6 +1915,50 @@ pub trait Itertools: Iterator {
     }
 
     // non-adaptor methods
+    /// Advances the iterator and returns the next items grouped in an array of
+    /// a specific size.
+    ///
+    /// If there are enough elements to be grouped in an array, then the array
+    /// is returned inside `Some`, otherwise `None` is returned.
+    ///
+    /// ```
+    /// use itertools::Itertools;
+    ///
+    /// let mut iter = 1..5;
+    ///
+    /// assert_eq!(Some([1, 2]), iter.next_array());
+    /// ```
+    fn next_array<T, const N: usize>(&mut self) -> Option<[T; N]>
+    where
+        Self: Sized + Iterator<Item = T>,
+    {
+        next_array::next_array(self)
+    }
+
+    /// Collects all items from the iterator into an array of a specific size.
+    ///
+    /// If the number of elements inside the iterator is **exactly** equal to
+    /// the array size, then the array is returned inside `Some`, otherwise
+    /// `None` is returned.
+    ///
+    /// ```
+    /// use itertools::Itertools;
+    ///
+    /// let iter = 1..3;
+    ///
+    /// if let Some([x, y]) = iter.collect_array() {
+    ///     assert_eq!([x, y], [1, 2])
+    /// } else {
+    ///     panic!("Expected two elements")
+    /// }
+    /// ```
+    fn collect_array<T, const N: usize>(mut self) -> Option<[T; N]>
+    where
+        Self: Sized + Iterator<Item = T>,
+    {
+        self.next_array().filter(|_| self.next().is_none())
+    }
+
     /// Advances the iterator and returns the next items grouped in a tuple of
     /// a specific size (up to 12).
     ///

src/next_array.rs

@@ -0,0 +1,64 @@
+use core::mem::MaybeUninit;
+use core::ptr;
+
+/// Helper struct to build up an array element by element.
+struct ArrayBuilder<T, const N: usize> {
+    arr: [MaybeUninit<T>; N], // Invariant: arr[..len] is initialized.
+    len: usize,               // Invariant: len <= N.
+}
+
+impl<T, const N: usize> ArrayBuilder<T, N> {
+    pub fn new() -> Self {
+        Self {
+            arr: [(); N].map(|_| MaybeUninit::uninit()),
+            len: 0,
+        }
+    }
+
+    pub fn push(&mut self, value: T) {
+        // We maintain the invariant here that arr[..len] is initialized.
+        // Indexing with self.len also ensures self.len < N, and thus <= N after
+        // the increment.
+        self.arr[self.len] = MaybeUninit::new(value);
+        self.len += 1;
+    }
+
+    pub fn take(&mut self) -> Option<[T; N]> {
+        if self.len == N {
+            // Take the array, resetting the length back to zero.
+            let arr = core::mem::replace(&mut self.arr, [(); N].map(|_| MaybeUninit::uninit()));
+            self.len = 0;
+
+            // SAFETY: we had len == N, so all elements in arr are initialized.
+            Some(unsafe { arr.map(|v| v.assume_init()) })
+        } else {
+            None
+        }
+    }
+}
+
+impl<T, const N: usize> Drop for ArrayBuilder<T, N> {
+    fn drop(&mut self) {
+        unsafe {
+            // SAFETY: arr[..len] is initialized, so must be dropped.
+            // First we create a pointer to this initialized slice, then drop
+            // that slice in-place. The cast from *mut MaybeUninit<T> to *mut T
+            // is always sound by the layout guarantees of MaybeUninit.
+            let ptr_to_first: *mut MaybeUninit<T> = self.arr.as_mut_ptr();
+            let ptr_to_slice = ptr::slice_from_raw_parts_mut(ptr_to_first.cast::<T>(), self.len);
+            ptr::drop_in_place(ptr_to_slice);
+        }
+    }
+}
+
+/// Equivalent to `it.next_array()`.
+pub fn next_array<I, T, const N: usize>(it: &mut I) -> Option<[T; N]>
+where
+    I: Iterator<Item = T>,
+{
+    let mut builder = ArrayBuilder::new();
+    for _ in 0..N {
+        builder.push(it.next()?);
+    }
+    builder.take()
+}

tests/test_core.rs

@@ -372,3 +372,28 @@ fn product1() {
     assert_eq!(v[1..3].iter().cloned().product1::<i32>(), Some(2));
     assert_eq!(v[1..5].iter().cloned().product1::<i32>(), Some(24));
 }
+
+#[test]
+fn next_array() {
+    let v = [1, 2, 3, 4, 5];
+    let mut iter = v.iter();
+    assert_eq!(iter.next_array(), Some([]));
+    assert_eq!(iter.next_array().map(|[&x, &y]| [x, y]), Some([1, 2]));
+    assert_eq!(iter.next_array().map(|[&x, &y]| [x, y]), Some([3, 4]));
+    assert_eq!(iter.next_array::<_, 2>(), None);
+}
+
+#[test]
+fn collect_array() {
+    let v = [1, 2];
+    let iter = v.iter().cloned();
+    assert_eq!(iter.collect_array(), Some([1, 2]));
+
+    let v = [1];
+    let iter = v.iter().cloned();
+    assert_eq!(iter.collect_array::<_, 2>(), None);
+
+    let v = [1, 2, 3];
+    let iter = v.iter().cloned();
+    assert_eq!(iter.collect_array::<_, 2>(), None);
+}