Add array_chunks

src/arrays.rs

@@ -0,0 +1,134 @@
+use alloc::vec::Vec;
+
+use crate::next_array::ArrayBuilder;
+
+macro_rules! const_assert_positive {
+    ($N: ty) => {
+        trait StaticAssert<const N: usize> {
+            const ASSERT: bool;
+        }
+
+        impl<const N: usize> StaticAssert<N> for () {
+            const ASSERT: bool = {
+                assert!(N > 0);
+                true
+            };
+        }
+
+        assert!(<() as StaticAssert<N>>::ASSERT);
+    };
+}
+
+/// An iterator that groups the items in arrays of const generic size `N`.
+///
+/// See [`.next_array()`](crate::Itertools::next_array) for details.
+#[derive(Debug, Clone)]
+pub struct Arrays<I: Iterator, const N: usize> {
+    iter: I,
+    partial: Vec<I::Item>,
+}
+
+impl<I: Iterator, const N: usize> Arrays<I, N> {
+    pub(crate) fn new(iter: I) -> Self {
+        const_assert_positive!(N);
+
+        // TODO should we use iter.fuse() instead? Otherwise remainder may behave strangely
+        Self {
+            iter,
+            partial: Vec::new(),
+        }
+    }
+
+    /// Returns an iterator that yields all the items that have
+    /// not been included in any of the arrays. Use this to access the
+    /// leftover elements if the total number of elements yielded by
+    /// the original iterator is not a multiple of `N`.
+    ///
+    /// If `self` is not exhausted (i.e. `next()` has not returned `None`)  
+    /// then the iterator returned by `remainder()` will also include
+    /// the elements that *would* have been included in the arrays
+    /// produced by `next()`.
+    ///
+    /// ```
+    /// use itertools::Itertools;
+    ///
+    /// let mut it = (1..9).arrays();
+    /// assert_eq!(Some([1, 2, 3]), it.next());
+    /// assert_eq!(Some([4, 5, 6]), it.next());
+    /// assert_eq!(None, it.next());
+    /// itertools::assert_equal(it.remainder(), [7,8]);
+    ///
+    /// let mut it = (1..9).arrays();
+    /// assert_eq!(Some([1, 2, 3]), it.next());
+    /// itertools::assert_equal(it.remainder(), 4..9);
+    /// ```
+    pub fn remainder(self) -> impl Iterator<Item = I::Item> {
+        self.partial.into_iter().chain(self.iter)
+    }
+}
+
+impl<I: Iterator, const N: usize> Iterator for Arrays<I, N> {
+    type Item = [I::Item; N];
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if !self.partial.is_empty() {
+            return None;
+        }
+        let mut builder = ArrayBuilder::new();
+        for _ in 0..N {
+            if let Some(item) = self.iter.next() {
+                builder.push(item);
+            } else {
+                break;
+            }
+        }
+        if let Some(array) = builder.take() {
+            Some(array)
+        } else {
+            self.partial = builder.into_vec();
+            None
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        if N == 0 {
+            (usize::MAX, None)
+        } else {
+            let (lo, hi) = self.iter.size_hint();
+            (lo / N, hi.map(|hi| hi / N))
+        }
+    }
+}
+
+impl<I: ExactSizeIterator, const N: usize> ExactSizeIterator for Arrays<I, N> {}
+
+#[cfg(test)]
+mod tests {
+    use crate::Itertools;
+
+    fn exact_size_helper(it: impl Iterator) {
+        let (lo, hi) = it.size_hint();
+        let count = it.count();
+        assert_eq!(lo, count);
+        assert_eq!(hi, Some(count));
+    }
+
+    #[test]
+    fn exact_size_not_divisible() {
+        let it = (0..10).array_chunks::<3>();
+        exact_size_helper(it);
+    }
+
+    #[test]
+    fn exact_size_after_next() {
+        let mut it = (0..10).array_chunks::<3>();
+        _ = it.next();
+        exact_size_helper(it);
+    }
+
+    #[test]
+    fn exact_size_divisible() {
+        let it = (0..10).array_chunks::<5>();
+        exact_size_helper(it);
+    }
+}

src/lib.rs

@@ -97,6 +97,8 @@ pub mod structs {
         TakeWhileRef, TupleCombinations, Update, WhileSome,
     };
     #[cfg(feature = "use_alloc")]
+    pub use crate::arrays::Arrays;
+    #[cfg(feature = "use_alloc")]
     pub use crate::combinations::{ArrayCombinations, Combinations};
     #[cfg(feature = "use_alloc")]
     pub use crate::combinations_with_replacement::CombinationsWithReplacement;
@@ -171,6 +173,8 @@ pub use crate::unziptuple::{multiunzip, MultiUnzip};
 pub use crate::with_position::Position;
 pub use crate::ziptuple::multizip;
 mod adaptors;
+#[cfg(feature = "use_alloc")]
+mod arrays;
 mod either_or_both;
 pub use crate::either_or_both::EitherOrBoth;
 #[doc(hidden)]
@@ -741,6 +745,55 @@ pub trait Itertools: Iterator {
         groupbylazy::new_chunks(self, size)
     }
 
+    /// Return an iterator that groups the items in arrays of const generic size `N`.
+    ///
+    /// Use the method `.remainder()` to access leftover items in case
+    /// the number of items yielded by the original iterator is not a multiple of `N`.
+    ///
+    /// `N == 0` is a compile-time (but post-monomorphization) error.
+    ///
+    /// See also the method [`.next_array()`](Itertools::next_array).
+    ///
+    /// ```rust
+    /// use itertools::Itertools;
+    /// let mut v = Vec::new();
+    /// for [a, b] in (1..5).arrays() {
+    ///     v.push([a, b]);
+    /// }
+    /// assert_eq!(v, vec![[1, 2], [3, 4]]);
+    ///
+    /// let mut it = (1..9).arrays();
+    /// assert_eq!(Some([1, 2, 3]), it.next());
+    /// assert_eq!(Some([4, 5, 6]), it.next());
+    /// assert_eq!(None, it.next());
+    /// itertools::assert_equal(it.remainder(), [7,8]);
+    ///
+    /// // this requires a type hint
+    /// let it = (1..7).arrays::<3>();
+    /// itertools::assert_equal(it, vec![[1, 2, 3], [4, 5, 6]]);
+    ///
+    /// // you can also specify the complete type
+    /// use itertools::Arrays;
+    /// use std::ops::Range;
+    ///
+    /// let it: Arrays<Range<u32>, 3> = (1..7).arrays();
+    /// itertools::assert_equal(it, vec![[1, 2, 3], [4, 5, 6]]);
+    /// ```
+    ///
+    /// ```compile_fail
+    /// use itertools::Itertools;
+    ///
+    /// let mut it = (1..5).arrays::<0>();
+    /// assert_eq!(Some([]), it.next());
+    /// ```
+    #[cfg(feature = "use_alloc")]
+    fn arrays<const N: usize>(self) -> Arrays<Self, N>
+    where
+        Self: Sized,
+    {
+        Arrays::new(self)
+    }
+
     /// Return an iterator over all contiguous windows producing tuples of
     /// a specific size (up to 12).
     ///

src/next_array.rs

@@ -1,7 +1,9 @@
+#[cfg(feature = "use_alloc")]
+use alloc::vec::Vec;
 use core::mem::{self, MaybeUninit};
 
 /// An array of at most `N` elements.
-struct ArrayBuilder<T, const N: usize> {
+pub(crate) struct ArrayBuilder<T, const N: usize> {
     /// The (possibly uninitialized) elements of the `ArrayBuilder`.
     ///
     /// # Safety
@@ -86,6 +88,23 @@ impl<T, const N: usize> ArrayBuilder<T, N> {
             None
         }
     }
+
+    #[cfg(feature = "use_alloc")]
+    pub(crate) fn into_vec(mut self) -> Vec<T> {
+        let len = self.len;
+        // SAFETY: Decreasing the value of `self.len` cannot violate the
+        // safety invariant on `self.arr`.
+        self.len = 0;
+        (0..len)
+            .map(|i| {
+                // SAFETY: Since `self.len` is 0, `self.arr` may safely contain
+                // uninitialized elements.
+                let item = mem::replace(&mut self.arr[i], MaybeUninit::uninit());
+                // SAFETY: we know that item is valid since i < len
+                unsafe { item.assume_init() }
+            })
+            .collect()
+    }
 }
 
 impl<T, const N: usize> AsMut<[T]> for ArrayBuilder<T, N> {