implement arrays and next_array

src/array_impl.rs

@@ -0,0 +1,172 @@
+/// An iterator that groups the items in arrays of a specific size.
+///
+/// See [`.arrays()`](crate::Itertools::arrays) for more information.
+#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
+pub struct Arrays<I: Iterator, const N: usize>
+{
+    iter: I,
+    buf: ArrayVec<I::Item, N>,
+}
+
+impl<I: Iterator, const N: usize> Arrays<I, N> {
+    /// Creates a new Arrays iterator.
+    /// See [`.arrays()`](crate::Itertools::arrays) for more information.
+    pub fn new(iter: I) -> Self {
+        Self { iter, buf: ArrayVec::new() }
+    }
+}
+impl<I: Iterator, const N: usize> Arrays<I, N> where I::Item: Clone {
+    /// Returns any remaining data left in the iterator.
+    /// This is useful if the iterator has left over values
+    /// that didn't fit into the array size.
+    pub fn remaining(&self) -> &[I::Item] {
+        self.buf.into_slice()
+    }
+}
+
+impl<I: Iterator, const N: usize> Iterator for Arrays<I, N> {
+    type Item = [I::Item; N];
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // SAFETY:
+        // ArrayVec::push_unchecked is safe as long as len < N.
+        // This is guaranteed by the for loop
+        unsafe {
+            for _ in self.buf.len()..N {
+                self.buf.push_unchecked(self.iter.next()?)
+            }
+
+            Some(self.buf.take_unchecked())
+        }
+    }
+}
+
+/// Reads N elements from the iter, returning them in an array. If there's not enough
+/// elements, returns None.
+pub fn next_array<I: Iterator, const N: usize>(iter: &mut I) -> Option<[I::Item; N]> {
+    let mut array_vec = ArrayVec::new();
+
+    unsafe {
+        // SAFETY:
+        // ArrayVec::push_unchecked is safe as long as len < N.
+        // This is guaranteed by the for loop
+        for _ in 0..N {
+            array_vec.push_unchecked(iter.next()?)
+        }
+
+        // SAFETY:
+        // We have guaranteed to have filled all N elements
+        Some(array_vec.into_array_unchecked())
+    }
+}
+
+// ArrayVec is a safe wrapper around a [T; N].
+// It allows safely initialising an empty array
+pub struct ArrayVec<T, const N: usize> {
+    data: std::mem::MaybeUninit<[T; N]>,
+    len: usize,
+}
+
+impl<T, const N: usize> Drop for ArrayVec<T, N> {
+    fn drop(&mut self) {
+        // SAFETY:
+        // The contract of the ArrayVec ensures that data[..len] is initialised
+        unsafe {
+            let ptr = self.data.as_mut_ptr() as *mut T;
+            drop(std::slice::from_raw_parts_mut(ptr, self.len));
+        }
+    }
+}
+
+impl<T, const N: usize> ArrayVec<T, N> {
+    /// Creates a new empty ArrayVec
+    pub const fn new() -> Self {
+        Self {
+            data: std::mem::MaybeUninit::uninit(),
+            len: 0,
+        }
+    }
+
+    /// Returns the number of initialised elements in the ArrayVec
+    pub fn len(&self) -> usize {
+        self.len
+    }
+
+    /// Returns whether the ArrayVec is empty
+    pub fn is_empty(&self) -> bool {
+        self.len == 0
+    }
+
+    /// Returns whether the ArrayVec is full
+    pub fn is_full(&self) -> bool {
+        self.len == N
+    }
+
+    /// Push a value into the ArrayVec
+    ///
+    /// Panics:
+    /// If the ArrayVec is full, this function will panic
+    pub fn push(&mut self, v: T) {
+        assert!(!self.is_full());
+        // SAFETY: asserted that self is not full
+        unsafe { self.push_unchecked(v) }
+    }
+
+    /// Push a value into the ArrayVec
+    ///
+    /// Unsafety:
+    /// The ArrayVec must not be full. If the ArrayVec is full, this function will try write data
+    /// out-of-bounds.
+    pub unsafe fn push_unchecked(&mut self, v: T) {
+        // The contract of ArrayVec guarantees that the value at self.len, if < N,
+        // is uninitialised, and therefore does not need dropping. So use write to
+        // overwrite the value
+        let ptr = (self.data.as_mut_ptr() as *mut T).add(self.len);
+        std::ptr::write(ptr, v);
+        self.len += 1;
+    }
+
+    /// If the ArrayVec is full, returns the data owned. Otherwise, returns None
+    pub fn into_array(self) -> Option<[T; N]> {
+        if self.is_full() {
+            // SAFETY:
+            // If len == N, then all the data is initialised and this is safe
+            unsafe { Some(self.into_array_unchecked()) }
+        } else {
+            None
+        }
+    }
+
+    /// Returns the data owned by the ArrayVec
+    ///
+    /// Unsafety:
+    /// The ArrayVec must be full. If it is not full, some of the values in the array will be
+    /// unintialised. This will cause undefined behaviour
+    unsafe fn into_array_unchecked(mut self) -> [T; N] {
+        // move out without dropping
+        let data = std::mem::replace(&mut self.data, std::mem::MaybeUninit::uninit().assume_init());
+        std::mem::forget(self);
+        data.assume_init()
+    }
+
+    /// Returns the data owned by the ArrayVec, resetting the ArrayVec to an empty state
+    ///
+    /// Unsafety:
+    /// The ArrayVec must be full. If it is not full, some of the values in the array will be
+    /// unintialised. This will cause undefined behaviour
+    unsafe fn take_unchecked(&mut self) -> [T; N] {
+        let data = std::mem::replace(&mut self.data, std::mem::MaybeUninit::uninit().assume_init());
+        self.len = 0;
+        data.assume_init()
+    }
+
+    /// Borrows the initialised data in the ArrayVec
+    pub fn into_slice(&self) -> &[T] {
+        unsafe {
+            // let data = std::mem::replace(&mut self.data, std::mem::MaybeUninit::uninit().assume_init());
+            // let len = self.len;
+            // std::mem::forget(self);
+            std::slice::from_raw_parts(self.data.as_ptr() as *const T, self.len)
+        }
+    }
+}

src/lib.rs

@@ -112,6 +112,7 @@ pub mod structs {
     pub use crate::adaptors::{MapResults, Step};
     #[cfg(feature = "use_alloc")]
     pub use crate::adaptors::MultiProduct;
+    pub use crate::array_impl::Arrays;
     #[cfg(feature = "use_alloc")]
     pub use crate::combinations::Combinations;
     #[cfg(feature = "use_alloc")]
@@ -181,6 +182,7 @@ pub use crate::sources::{repeat_call, unfold, iterate};
 pub use crate::with_position::Position;
 pub use crate::ziptuple::multizip;
 mod adaptors;
+mod array_impl;
 mod either_or_both;
 pub use crate::either_or_both::EitherOrBoth;
 #[doc(hidden)]
@@ -775,6 +777,42 @@ pub trait Itertools : Iterator {
         tuple_impl::tuples(self)
     }
 
+    /// Return an iterator that groups the items in arrays of a specific size
+    ///
+    /// See also the method [`.next_array()`](#method.next_array).
+    ///
+    /// ```
+    /// 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..8).arrays();
+    /// assert_eq!(Some([1, 2, 3]), it.next());
+    /// assert_eq!(Some([4, 5, 6]), it.next());
+    /// assert_eq!(None, it.next());
+    /// // get any remaining data
+    /// assert_eq!(vec![7], it.remaining());
+    ///
+    /// // 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]]);
+    /// ```
+    fn arrays<const N: usize>(self) -> array_impl::Arrays<Self, N>
+        where Self: Sized + Iterator,
+    {
+        array_impl::Arrays::new(self)
+    }
+
     /// Split into an iterator pair that both yield all elements from
     /// the original iterator.
     ///
@@ -1702,6 +1740,25 @@ pub trait Itertools : Iterator {
         T::collect_from_iter_no_buf(self)
     }
 
+    /// 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 a tuple, then the tuple 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<const N: usize>(&mut self) -> Option<[Self::Item; N]>
+        where Self: Sized + Iterator
+    {
+        array_impl::next_array(self)
+    }
+
     /// Collects all items from the iterator into a tuple of a specific size
     /// (up to 12).
     ///

tests/arrays.rs

@@ -0,0 +1,39 @@
+use itertools::Itertools;
+
+#[test]
+fn arrays() {
+    let v = [1, 2, 3, 4, 5];
+    let mut iter = v.iter().cloned().arrays();
+    assert_eq!(Some([1]), iter.next());
+    assert_eq!(Some([2]), iter.next());
+    assert_eq!(Some([3]), iter.next());
+    assert_eq!(Some([4]), iter.next());
+    assert_eq!(Some([5]), iter.next());
+    assert_eq!(None, iter.next());
+    assert_eq!(Vec::<i32>::new(), iter.remaining());
+
+    let mut iter = v.iter().cloned().arrays();
+    assert_eq!(Some([1, 2]), iter.next());
+    assert_eq!(Some([3, 4]), iter.next());
+    assert_eq!(None, iter.next());
+    assert_eq!(vec![5], iter.remaining());
+
+    let mut iter = v.iter().cloned().arrays();
+    assert_eq!(Some([1, 2, 3]), iter.next());
+    assert_eq!(None, iter.next());
+    assert_eq!(vec![4, 5], iter.remaining());
+
+    let mut iter = v.iter().cloned().arrays();
+    assert_eq!(Some([1, 2, 3, 4]), iter.next());
+    assert_eq!(None, iter.next());
+    assert_eq!(vec![5], iter.remaining());
+}
+
+#[test]
+fn next_array() {
+    let v = [1, 2, 3, 4, 5];
+    let mut iter = v.iter();
+    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);
+}

tests/quick.rs

@@ -1073,6 +1073,32 @@ quickcheck! {
     }
 }
 
+quickcheck! {
+    fn equal_arrays_1(a: Vec<u8>) -> bool {
+        let x = a.chunks(1).map(|s| [&s[0]]);
+        let y = a.iter().arrays::<1>();
+        itertools::equal(x, y)
+    }
+
+    fn equal_arrays_2(a: Vec<u8>) -> bool {
+        let x = a.chunks(2).filter(|s| s.len() == 2).map(|s| [&s[0], &s[1]]);
+        let y = a.iter().arrays::<2>();
+        itertools::equal(x, y)
+    }
+
+    fn equal_arrays_3(a: Vec<u8>) -> bool {
+        let x = a.chunks(3).filter(|s| s.len() == 3).map(|s| [&s[0], &s[1], &s[2]]);
+        let y = a.iter().arrays::<3>();
+        itertools::equal(x, y)
+    }
+
+    fn equal_arrays_4(a: Vec<u8>) -> bool {
+        let x = a.chunks(4).filter(|s| s.len() == 4).map(|s| [&s[0], &s[1], &s[2], &s[3]]);
+        let y = a.iter().arrays::<4>();
+        itertools::equal(x, y)
+    }
+}
+
 // with_position
 quickcheck! {
     fn with_position_exact_size_1(a: Vec<u8>) -> bool {