diff --git a/src/arrays.rs b/src/arrays.rs new file mode 100644 index 000000000..c955c5f86 --- /dev/null +++ b/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 ASSERT: bool; + } + + impl StaticAssert for () { + const ASSERT: bool = { + assert!(N > 0); + true + }; + } + + assert!(<() as StaticAssert>::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 { + iter: I, + partial: Vec, +} + +impl Arrays { + 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 { + self.partial.into_iter().chain(self.iter) + } +} + +impl Iterator for Arrays { + type Item = [I::Item; N]; + + fn next(&mut self) -> Option { + 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) { + if N == 0 { + (usize::MAX, None) + } else { + let (lo, hi) = self.iter.size_hint(); + (lo / N, hi.map(|hi| hi / N)) + } + } +} + +impl ExactSizeIterator for Arrays {} + +#[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); + } +} diff --git a/src/lib.rs b/src/lib.rs index ff117800a..a28404934 100644 --- a/src/lib.rs +++ b/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, 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(self) -> Arrays + where + Self: Sized, + { + Arrays::new(self) + } + /// Return an iterator over all contiguous windows producing tuples of /// a specific size (up to 12). /// diff --git a/src/next_array.rs b/src/next_array.rs index 86480b197..09e27fc6c 100644 --- a/src/next_array.rs +++ b/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 { +pub(crate) struct ArrayBuilder { /// The (possibly uninitialized) elements of the `ArrayBuilder`. /// /// # Safety @@ -86,6 +88,23 @@ impl ArrayBuilder { None } } + + #[cfg(feature = "use_alloc")] + pub(crate) fn into_vec(mut self) -> Vec { + 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 AsMut<[T]> for ArrayBuilder {