Unify chunks, fold_chunks, and fold_chunks_with

rayon-rs · Nov 16, 2022 · a119f23 · a119f23
1 parent 911d6d0
commit a119f23
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 310 deletions.
diff --git a/src/iter/chunks.rs b/src/iter/chunks.rs
@@ -90,38 +90,46 @@ where
             where
                 P: Producer<Item = T>,
             {
-                self.callback.callback(ChunkProducer {
-                    chunk_size: self.size,
-                    len: self.len,
-                    base,
-                })
+                let producer = ChunkProducer::new(self.size, self.len, base, Vec::from_iter);
+                self.callback.callback(producer)
             }
         }
     }
 }
 
-struct ChunkProducer<P>
-where
-    P: Producer,
-{
+pub(super) struct ChunkProducer<P, F> {
     chunk_size: usize,
     len: usize,
     base: P,
+    map: F,
 }
 
-impl<P> Producer for ChunkProducer<P>
+impl<P, F> ChunkProducer<P, F> {
+    pub(super) fn new(chunk_size: usize, len: usize, base: P, map: F) -> Self {
+        Self {
+            chunk_size,
+            len,
+            base,
+            map,
+        }
+    }
+}
+
+impl<P, F, T> Producer for ChunkProducer<P, F>
 where
     P: Producer,
+    F: Fn(P::IntoIter) -> T + Send + Clone,
 {
-    type Item = Vec<P::Item>;
-    type IntoIter = ChunkSeq<P>;
+    type Item = T;
+    type IntoIter = std::iter::Map<ChunkSeq<P>, F>;
 
     fn into_iter(self) -> Self::IntoIter {
-        ChunkSeq {
+        let chunks = ChunkSeq {
             chunk_size: self.chunk_size,
             len: self.len,
             inner: if self.len > 0 { Some(self.base) } else { None },
-        }
+        };
+        chunks.map(self.map)
     }
 
     fn split_at(self, index: usize) -> (Self, Self) {
@@ -132,11 +140,13 @@ where
                 chunk_size: self.chunk_size,
                 len: elem_index,
                 base: left,
+                map: self.map.clone(),
             },
             ChunkProducer {
                 chunk_size: self.chunk_size,
                 len: self.len - elem_index,
                 base: right,
+                map: self.map,
             },
         )
     }
@@ -150,7 +160,7 @@ where
     }
 }
 
-struct ChunkSeq<P> {
+pub(super) struct ChunkSeq<P> {
     chunk_size: usize,
     len: usize,
     inner: Option<P>,
@@ -160,19 +170,19 @@ impl<P> Iterator for ChunkSeq<P>
 where
     P: Producer,
 {
-    type Item = Vec<P::Item>;
+    type Item = P::IntoIter;
 
     fn next(&mut self) -> Option<Self::Item> {
         let producer = self.inner.take()?;
         if self.len > self.chunk_size {
             let (left, right) = producer.split_at(self.chunk_size);
             self.inner = Some(right);
             self.len -= self.chunk_size;
-            Some(left.into_iter().collect())
+            Some(left.into_iter())
         } else {
             debug_assert!(self.len > 0);
             self.len = 0;
-            Some(producer.into_iter().collect())
+            Some(producer.into_iter())
         }
     }
 
@@ -206,11 +216,11 @@ where
             let (left, right) = producer.split_at(self.len - size);
             self.inner = Some(left);
             self.len -= size;
-            Some(right.into_iter().collect())
+            Some(right.into_iter())
         } else {
             debug_assert!(self.len > 0);
             self.len = 0;
-            Some(producer.into_iter().collect())
+            Some(producer.into_iter())
         }
     }
 }
diff --git a/src/iter/fold_chunks.rs b/src/iter/fold_chunks.rs
@@ -1,6 +1,6 @@
-use std::cmp::min;
 use std::fmt::{self, Debug};
 
+use super::chunks::ChunkProducer;
 use super::plumbing::*;
 use super::*;
 use crate::math::div_round_up;
@@ -123,155 +123,16 @@ where
             where
                 P: Producer<Item = T>,
             {
-                self.callback.callback(FoldChunksProducer {
-                    chunk_size: self.chunk_size,
-                    len: self.len,
-                    identity: &self.identity,
-                    fold_op: &self.fold_op,
-                    base,
-                })
+                let identity = &self.identity;
+                let fold_op = &self.fold_op;
+                let fold_iter = move |iter: P::IntoIter| iter.fold(identity(), fold_op);
+                let producer = ChunkProducer::new(self.chunk_size, self.len, base, fold_iter);
+                self.callback.callback(producer)
             }
         }
     }
 }
 
-struct FoldChunksProducer<'f, P, ID, F>
-where
-    P: Producer,
-{
-    chunk_size: usize,
-    len: usize,
-    identity: &'f ID,
-    fold_op: &'f F,
-    base: P,
-}
-
-impl<'f, P, ID, U, F> Producer for FoldChunksProducer<'f, P, ID, F>
-where
-    P: Producer,
-    ID: Fn() -> U + Send + Sync,
-    F: Fn(U, P::Item) -> U + Send + Sync,
-{
-    type Item = F::Output;
-    type IntoIter = FoldChunksSeq<'f, P, ID, F>;
-
-    fn into_iter(self) -> Self::IntoIter {
-        FoldChunksSeq {
-            chunk_size: self.chunk_size,
-            len: self.len,
-            identity: self.identity,
-            fold_op: self.fold_op,
-            inner: if self.len > 0 { Some(self.base) } else { None },
-        }
-    }
-
-    fn min_len(&self) -> usize {
-        div_round_up(self.base.min_len(), self.chunk_size)
-    }
-
-    fn max_len(&self) -> usize {
-        self.base.max_len() / self.chunk_size
-    }
-
-    fn split_at(self, index: usize) -> (Self, Self) {
-        let elem_index = min(index * self.chunk_size, self.len);
-        let (left, right) = self.base.split_at(elem_index);
-        (
-            FoldChunksProducer {
-                chunk_size: self.chunk_size,
-                len: elem_index,
-                identity: self.identity,
-                fold_op: self.fold_op,
-                base: left,
-            },
-            FoldChunksProducer {
-                chunk_size: self.chunk_size,
-                len: self.len - elem_index,
-                identity: self.identity,
-                fold_op: self.fold_op,
-                base: right,
-            },
-        )
-    }
-}
-
-struct FoldChunksSeq<'f, P, ID, F> {
-    chunk_size: usize,
-    len: usize,
-    identity: &'f ID,
-    fold_op: &'f F,
-    inner: Option<P>,
-}
-
-impl<'f, P, ID, U, F> Iterator for FoldChunksSeq<'f, P, ID, F>
-where
-    P: Producer,
-    ID: Fn() -> U + Send + Sync,
-    F: Fn(U, P::Item) -> U + Send + Sync,
-{
-    type Item = U;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        let producer = self.inner.take()?;
-        if self.len > self.chunk_size {
-            let (left, right) = producer.split_at(self.chunk_size);
-            self.inner = Some(right);
-            self.len -= self.chunk_size;
-            let chunk = left.into_iter();
-            Some(chunk.fold((self.identity)(), self.fold_op))
-        } else {
-            debug_assert!(self.len > 0);
-            self.len = 0;
-            let chunk = producer.into_iter();
-            Some(chunk.fold((self.identity)(), self.fold_op))
-        }
-    }
-
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        let len = self.len();
-        (len, Some(len))
-    }
-}
-
-impl<'f, P, ID, U, F> ExactSizeIterator for FoldChunksSeq<'f, P, ID, F>
-where
-    P: Producer,
-    ID: Fn() -> U + Send + Sync,
-    F: Fn(U, P::Item) -> U + Send + Sync,
-{
-    #[inline]
-    fn len(&self) -> usize {
-        div_round_up(self.len, self.chunk_size)
-    }
-}
-
-impl<'f, P, ID, U, F> DoubleEndedIterator for FoldChunksSeq<'f, P, ID, F>
-where
-    P: Producer,
-    ID: Fn() -> U + Send + Sync,
-    F: Fn(U, P::Item) -> U + Send + Sync,
-{
-    fn next_back(&mut self) -> Option<Self::Item> {
-        let producer = self.inner.take()?;
-        if self.len > self.chunk_size {
-            let mut size = self.len % self.chunk_size;
-            if size == 0 {
-                size = self.chunk_size;
-            }
-            let (left, right) = producer.split_at(self.len - size);
-            self.inner = Some(left);
-            self.len -= size;
-            let chunk = right.into_iter();
-            Some(chunk.fold((self.identity)(), self.fold_op))
-        } else {
-            debug_assert!(self.len > 0);
-            self.len = 0;
-            let chunk = producer.into_iter();
-            Some(chunk.fold((self.identity)(), self.fold_op))
-        }
-    }
-}
-
 #[cfg(test)]
 mod test {
     use super::*;