Add next_array and collect_array

Cargo.toml

@@ -25,6 +25,9 @@ test = false
 [dependencies]
 either = { version = "1.0", default-features = false }
 
+[build-dependencies]
+version_check = "0.9"
+
 [dev-dependencies]
 rand = "0.7"
 criterion = "=0" # TODO how could this work with our minimum supported Rust version?

build.rs

@@ -0,0 +1,13 @@
+fn main() {
+    let is_nightly = version_check::is_feature_flaggable() == Some(true);
+    let is_at_least_1_34 = version_check::is_min_version("1.34.0").unwrap_or(false);
+    let is_at_least_1_51 = version_check::is_min_version("1.51.0").unwrap_or(false);
+
+    if !is_at_least_1_34 && !is_nightly {
+        println!("cargo:warning=itertools requires rustc => 1.34.0");
+    }
+
+    if is_at_least_1_51 || is_nightly {
+        println!("cargo:rustc-cfg=has_min_const_generics");
+    }
+}

src/lib.rs

@@ -42,7 +42,7 @@
 //!
 //! ## Rust Version
 //!
-//! This version of itertools requires Rust 1.32 or later.
+//! This version of itertools requires Rust 1.34 or later.
 #![doc(html_root_url="https://docs.rs/itertools/0.8/")]
 
 #[cfg(not(feature = "use_std"))]
@@ -214,6 +214,8 @@ mod merge_join;
 mod minmax;
 #[cfg(feature = "use_alloc")]
 mod multipeek_impl;
+#[cfg(has_min_const_generics)]
+mod next_array;
 mod pad_tail;
 #[cfg(feature = "use_alloc")]
 mod peek_nth;
@@ -1679,6 +1681,58 @@ 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.
+    ///
+    /// Requires Rust version 1.51 or later.
+    ///
+    /// ```
+    /// use itertools::Itertools;
+    ///
+    /// let mut iter = 1..5;
+    ///
+    /// assert_eq!(Some([1, 2]), iter.next_array());
+    /// ```
+    #[cfg(has_min_const_generics)]
+    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.
+    ///
+    /// Requires Rust version 1.51 or later.
+    ///
+    /// ```
+    /// 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")
+    /// }
+    /// ```
+    #[cfg(has_min_const_generics)]
+    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).
     ///
@@ -1699,6 +1753,7 @@ pub trait Itertools : Iterator {
         T::collect_from_iter_no_buf(self)
     }
 
+
     /// Collects all items from the iterator into a tuple of a specific size
     /// (up to 12).
     ///
@@ -1721,13 +1776,7 @@ pub trait Itertools : Iterator {
         where Self: Sized + Iterator<Item = T::Item>,
               T: traits::HomogeneousTuple
     {
-        match self.next_tuple() {
-            elt @ Some(_) => match self.next() {
-                Some(_) => None,
-                None => elt,
-            },
-            _ => None
-        }
+        self.next_tuple().filter(|_| self.next().is_none())
     }
 
 

src/next_array.rs

@@ -0,0 +1,80 @@
+use core::mem::MaybeUninit;
+
+/// Helper struct to build up an array element by element.
+struct ArrayBuilder<T, const N: usize> {
+    arr: [MaybeUninit<T>; N],
+    i: usize
+}
+
+impl<T, const N: usize> ArrayBuilder<T, N> {
+    pub fn new() -> Self {
+        Self { arr: maybe_uninit::uninit_array(), i: 0 }
+    }
+
+    pub unsafe fn push_unchecked(&mut self, x: T) {
+        debug_assert!(self.i < N);
+        *self.arr.get_unchecked_mut(self.i) = MaybeUninit::new(x);
+        self.i += 1;
+    }
+
+    pub fn take(mut self) -> Option<[T; N]> {
+        if self.i == N {
+            unsafe {
+                // SAFETY: prevent double drop.
+                self.i = 0;
+                // SAFETY: [MaybeUninit<T>; N] and [T; N] have the same layout.
+                let init_arr_ptr = &self.arr as *const _ as *const [T; N];
+                Some(core::ptr::read(init_arr_ptr))
+            }
+        } else {
+            None
+        }
+    }
+}
+
+impl<T, const N: usize> Drop for ArrayBuilder<T, N> {
+    fn drop(&mut self) {
+        unsafe {
+            // SAFETY: we only loop over the initialized portion.
+            for el in &mut self.arr[..self.i] {
+                maybe_uninit::assume_init_drop(el)
+            }
+        }
+    }
+}
+
+
+
+/// 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 el in it.take(N) {
+        unsafe {
+            // SAFETY: the take(N) guarantees we never go out of bounds.
+            builder.push_unchecked(el);
+        }
+    }
+    builder.take()
+}
+
+
+
+/// Replacements for unstable core methods, copied from stdlib.
+mod maybe_uninit {
+    use core::mem::MaybeUninit;
+
+    pub fn uninit_array<T, const N: usize>() -> [MaybeUninit<T>; N] {
+        // SAFETY: an uninitialized `[MaybeUninit<_>; N]` is valid.
+        unsafe { MaybeUninit::<[MaybeUninit<T>; N]>::uninit().assume_init() }
+    }
+
+    pub unsafe fn assume_init_drop<T>(u: &mut MaybeUninit<T>) {
+        // SAFETY: the caller must guarantee that `self` is initialized and
+        // satisfies all invariants of `T`.
+        // Dropping the value in place is safe if that is the case.
+        core::ptr::drop_in_place(u.as_mut_ptr())
+    }
+}

tests/test_core.rs

@@ -301,3 +301,29 @@ 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);
+}