Implemeneted 'chunks' operation

asparit/src/core/iterator.rs

@@ -8,6 +8,7 @@ use super::{
 use crate::{
     iter::{
         chain::Chain,
+        chunks::Chunks,
         cloned::Cloned,
         collect::Collect,
         copied::Copied,
@@ -1889,6 +1890,32 @@ pub trait IndexedParallelIterator<'a>: ParallelIterator<'a> {
         }
     }
 
+    /// Splits an iterator up into fixed-size chunks.
+    ///
+    /// Returns an iterator that returns `Vec`s of the given number of elements.
+    /// If the number of elements in the iterator is not divisible by `chunk_size`,
+    /// the last chunk may be shorter than `chunk_size`.
+    ///
+    /// See also [`par_chunks()`] and [`par_chunks_mut()`] for similar behavior on
+    /// slices, without having to allocate intermediate `Vec`s for the chunks.
+    ///
+    /// [`par_chunks()`]: ../slice/trait.ParallelSlice.html#method.par_chunks
+    /// [`par_chunks_mut()`]: ../slice/trait.ParallelSliceMut.html#method.par_chunks_mut
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use rayon::prelude::*;
+    /// let a = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+    /// let r: Vec<Vec<i32>> = a.into_par_iter().chunks(3).collect();
+    /// assert_eq!(r, vec![vec![1,2,3], vec![4,5,6], vec![7,8,9], vec![10]]);
+    /// ```
+    fn chunks(self, chunk_size: usize) -> Chunks<Self> {
+        assert!(chunk_size != 0, "chunk_size must not be zero");
+
+        Chunks::new(self, chunk_size)
+    }
+
     /// Creates an iterator that yields the first `n` elements.
     ///
     /// # Examples

asparit/src/iter/chunks.rs

@@ -0,0 +1,250 @@
+use std::cmp::min;
+use std::iter::{DoubleEndedIterator, ExactSizeIterator, Iterator};
+
+use crate::{
+    Consumer, Executor, ExecutorCallback, IndexedParallelIterator, IndexedProducer,
+    IndexedProducerCallback, ParallelIterator, Reducer, Setup, WithIndexedProducer, WithSetup,
+};
+
+pub struct Chunks<X> {
+    base: X,
+    size: usize,
+}
+
+impl<X> Chunks<X> {
+    pub fn new(base: X, size: usize) -> Self {
+        Self { base, size }
+    }
+}
+
+impl<'a, X, I> ParallelIterator<'a> for Chunks<X>
+where
+    X: IndexedParallelIterator<'a, Item = I> + WithIndexedProducer<'a, Item = I>,
+    I: Send + 'a,
+{
+    type Item = Vec<I>;
+
+    fn drive<E, C, D, R>(self, executor: E, consumer: C) -> E::Result
+    where
+        E: Executor<'a, D>,
+        C: Consumer<Self::Item, Result = D, Reducer = R> + 'a,
+        D: Send + 'a,
+        R: Reducer<D> + Send + 'a,
+    {
+        self.with_indexed_producer(ExecutorCallback::new(executor, consumer))
+    }
+
+    fn len_hint_opt(&self) -> Option<usize> {
+        Some(self.len_hint())
+    }
+}
+
+impl<'a, X, I> IndexedParallelIterator<'a> for Chunks<X>
+where
+    X: IndexedParallelIterator<'a, Item = I> + WithIndexedProducer<'a, Item = I>,
+    I: Send + 'a,
+{
+    fn drive_indexed<E, C, D, R>(self, executor: E, consumer: C) -> E::Result
+    where
+        E: Executor<'a, D>,
+        C: Consumer<Self::Item, Result = D, Reducer = R> + 'a,
+        D: Send + 'a,
+        R: Reducer<D> + Send + 'a,
+    {
+        self.with_indexed_producer(ExecutorCallback::new(executor, consumer))
+    }
+
+    fn len_hint(&self) -> usize {
+        let mut len = self.base.len_hint();
+
+        if len > 0 {
+            len = (len - 1) / self.size - 1;
+        }
+
+        len
+    }
+}
+
+impl<'a, X> WithIndexedProducer<'a> for Chunks<X>
+where
+    X: WithIndexedProducer<'a>,
+{
+    type Item = Vec<X::Item>;
+
+    fn with_indexed_producer<CB>(self, base: CB) -> CB::Output
+    where
+        CB: IndexedProducerCallback<'a, Self::Item>,
+    {
+        self.base.with_indexed_producer(ChunksCallback {
+            base,
+            size: self.size,
+        })
+    }
+}
+
+/* ChunksCallback */
+
+struct ChunksCallback<CB> {
+    base: CB,
+    size: usize,
+}
+
+impl<'a, CB, I> IndexedProducerCallback<'a, I> for ChunksCallback<CB>
+where
+    CB: IndexedProducerCallback<'a, Vec<I>>,
+{
+    type Output = CB::Output;
+
+    fn callback<P>(self, base: P) -> CB::Output
+    where
+        P: IndexedProducer<Item = I> + 'a,
+    {
+        self.base.callback(ChunkProducer {
+            base,
+            size: self.size,
+        })
+    }
+}
+
+/* ChunkProducer */
+
+struct ChunkProducer<P> {
+    base: P,
+    size: usize,
+}
+
+impl<P> WithSetup for ChunkProducer<P>
+where
+    P: WithSetup,
+{
+    fn setup(&self) -> Setup {
+        self.base.setup()
+    }
+}
+
+impl<P> IndexedProducer for ChunkProducer<P>
+where
+    P: IndexedProducer,
+{
+    type Item = Vec<P::Item>;
+    type IntoIter = ChunksIter<P>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        ChunksIter {
+            producer: if self.len() > 0 {
+                Some(self.base)
+            } else {
+                None
+            },
+            size: self.size,
+        }
+    }
+
+    fn len(&self) -> usize {
+        let len = self.base.len();
+
+        if len > 0 {
+            (len - 1) / self.size + 1
+        } else {
+            0
+        }
+    }
+
+    fn split_at(self, index: usize) -> (Self, Self) {
+        let index = min(index * self.size, self.base.len());
+        let (left, right) = self.base.split_at(index);
+
+        let left = Self {
+            base: left,
+            size: self.size,
+        };
+        let right = Self {
+            base: right,
+            size: self.size,
+        };
+
+        (left, right)
+    }
+}
+
+/* ChunksIter */
+
+struct ChunksIter<P> {
+    producer: Option<P>,
+    size: usize,
+}
+
+impl<P> Iterator for ChunksIter<P>
+where
+    P: IndexedProducer,
+{
+    type Item = Vec<P::Item>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let producer = self.producer.take()?;
+        let producer = if producer.len() > self.size {
+            let index = self.size;
+            let (left, right) = producer.split_at(index);
+
+            self.producer = Some(right);
+
+            left
+        } else {
+            producer
+        };
+
+        Some(producer.into_iter().collect())
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = self.len();
+
+        (len, Some(len))
+    }
+}
+
+impl<P> DoubleEndedIterator for ChunksIter<P>
+where
+    P: IndexedProducer,
+{
+    fn next_back(&mut self) -> Option<Self::Item> {
+        let producer = self.producer.take()?;
+        let producer = if producer.len() > self.size {
+            let mut size = producer.len() % self.size;
+            if size == 0 {
+                size = self.size;
+            }
+
+            let index = producer.len() - size;
+            let (left, right) = producer.split_at(index);
+
+            self.producer = Some(left);
+
+            right
+        } else {
+            producer
+        };
+
+        Some(producer.into_iter().collect())
+    }
+}
+
+impl<P> ExactSizeIterator for ChunksIter<P>
+where
+    P: IndexedProducer,
+{
+    #[inline]
+    fn len(&self) -> usize {
+        let len = self
+            .producer
+            .as_ref()
+            .map(IndexedProducer::len)
+            .unwrap_or_default();
+
+        if len > 0 {
+            (len - 1) / self.size + 1
+        } else {
+            0
+        }
+    }
+}

asparit/src/iter/mod.rs

@@ -1,4 +1,5 @@
 pub mod chain;
+pub mod chunks;
 pub mod cloned;
 pub mod collect;
 pub mod copied;
@@ -39,26 +40,10 @@ mod tests {
 
     #[tokio::test(flavor = "multi_thread")]
     async fn test_for_each() {
-        let a = vec![
-            vec![1usize, 2usize],
-            vec![3usize, 4usize],
-            vec![5usize, 6usize],
-        ];
-        let b = vec![
-            vec![7usize, 8usize],
-            vec![9usize, 10usize],
-            vec![11usize, 12usize],
-        ];
-
-        a.par_iter()
-            .cloned()
-            .chain(b)
+        vec![0usize, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+            .into_par_iter()
             .with_splits(1)
-            .interleave_shortest(
-                vec![vec![50, 51], vec![52, 53], vec![54, 55]]
-                    .into_par_iter()
-                    .take(2),
-            )
+            .chunks(4)
             .for_each(|x| {
                 dbg!(x);
             })