feat: arrays and next_array

creates an ArrayVec type that handles uninitialised array data

Author: conradludgate

src/array_impl.rs

@@ -0,0 +1,144 @@
+/// 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> {
+    pub fn new(iter: I) -> Self {
+        Self { iter, buf: ArrayVec::new() }
+    }
+}
+impl<I: Iterator, const N: usize> Arrays<I, N> where I::Item: Clone {
+    pub fn remaining(self) -> Vec<I::Item> {
+        self.buf.into_vec()
+    }
+}
+
+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())
+        }
+    }
+}
+
+pub fn next_array<I: Iterator, const N: usize>(iter: &mut I) -> Option<[I::Item; N]> {
+    let mut array_vec = ArrayVec::new();
+
+    // SAFETY:
+    // ArrayVec::push_unchecked is safe as long as len < N.
+    // This is guaranteed by the for loop
+    unsafe {
+        for _ in 0..N {
+            array_vec.push_unchecked(iter.next()?)
+        }
+    }
+
+    array_vec.into_array()
+}
+
+// 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> {
+    pub const fn new() -> Self {
+        Self {
+            data: std::mem::MaybeUninit::uninit(),
+            len: 0,
+        }
+    }
+
+    pub fn push(&mut self, v: T) {
+        assert!(self.len < N);
+        // SAFETY: asserts that len < N.
+        unsafe { self.push_unchecked(v) }
+    }
+
+    // Unsafety:
+    // len must be less than N. If `len < N` is guaranteed, this operation is safe
+    // This is because the contract of ArrayVec guarantees that if len < N, then the value
+    // at len is valid and uninitialised.
+    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;
+    }
+
+    pub fn len(&self) -> usize {
+        self.len
+    }
+
+    pub fn into_array(self) -> Option<[T; N]> {
+        if self.len == N {
+            // SAFETY:
+            // If len == N, then all the data is initialised and this is safe
+            unsafe { Some(self.into_array_unchecked()) }
+        } else {
+            None
+        }
+    }
+
+    // Unsafety:
+    // len must be equal to N. If `len == N` is guaranteed, this operation is safe.
+    // This is because the contract of ArrayVec guarantees that if len == N, all the values
+    // have been initialised correctly.
+    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()
+    }
+
+    // Unsafety:
+    // len must be equal to N. If `len == N` is guaranteed, this operation is safe.
+    // This is because the contract of ArrayVec guarantees that if len == N, all the values
+    // have been initialised correctly.
+    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()
+    }
+}
+
+impl<T: Clone, const N: usize> ArrayVec<T, N> {
+    pub fn into_vec(mut self) -> Vec<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(data.as_ptr() as *const T, len).to_vec()
+        }
+    }
+}

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 {