rayon/iter/
chunks.rs

1use super::plumbing::*;
2use super::*;
3use crate::math::div_round_up;
4
5/// `Chunks` is an iterator that groups elements of an underlying iterator.
6///
7/// This struct is created by the [`chunks()`] method on [`IndexedParallelIterator`]
8///
9/// [`chunks()`]: trait.IndexedParallelIterator.html#method.chunks
10/// [`IndexedParallelIterator`]: trait.IndexedParallelIterator.html
11#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
12#[derive(Debug, Clone)]
13pub struct Chunks<I>
14where
15    I: IndexedParallelIterator,
16{
17    size: usize,
18    i: I,
19}
20
21impl<I> Chunks<I>
22where
23    I: IndexedParallelIterator,
24{
25    /// Creates a new `Chunks` iterator
26    pub(super) fn new(i: I, size: usize) -> Self {
27        Chunks { i, size }
28    }
29}
30
31impl<I> ParallelIterator for Chunks<I>
32where
33    I: IndexedParallelIterator,
34{
35    type Item = Vec<I::Item>;
36
37    fn drive_unindexed<C>(self, consumer: C) -> C::Result
38    where
39        C: Consumer<Vec<I::Item>>,
40    {
41        bridge(self, consumer)
42    }
43
44    fn opt_len(&self) -> Option<usize> {
45        Some(self.len())
46    }
47}
48
49impl<I> IndexedParallelIterator for Chunks<I>
50where
51    I: IndexedParallelIterator,
52{
53    fn drive<C>(self, consumer: C) -> C::Result
54    where
55        C: Consumer<Self::Item>,
56    {
57        bridge(self, consumer)
58    }
59
60    fn len(&self) -> usize {
61        div_round_up(self.i.len(), self.size)
62    }
63
64    fn with_producer<CB>(self, callback: CB) -> CB::Output
65    where
66        CB: ProducerCallback<Self::Item>,
67    {
68        let len = self.i.len();
69        return self.i.with_producer(Callback {
70            size: self.size,
71            len,
72            callback,
73        });
74
75        struct Callback<CB> {
76            size: usize,
77            len: usize,
78            callback: CB,
79        }
80
81        impl<T, CB> ProducerCallback<T> for Callback<CB>
82        where
83            CB: ProducerCallback<Vec<T>>,
84        {
85            type Output = CB::Output;
86
87            fn callback<P>(self, base: P) -> CB::Output
88            where
89                P: Producer<Item = T>,
90            {
91                let producer = ChunkProducer::new(self.size, self.len, base, Vec::from_iter);
92                self.callback.callback(producer)
93            }
94        }
95    }
96}
97
98pub(super) struct ChunkProducer<P, F> {
99    chunk_size: usize,
100    len: usize,
101    base: P,
102    map: F,
103}
104
105impl<P, F> ChunkProducer<P, F> {
106    pub(super) fn new(chunk_size: usize, len: usize, base: P, map: F) -> Self {
107        Self {
108            chunk_size,
109            len,
110            base,
111            map,
112        }
113    }
114}
115
116impl<P, F, T> Producer for ChunkProducer<P, F>
117where
118    P: Producer,
119    F: Fn(P::IntoIter) -> T + Send + Clone,
120{
121    type Item = T;
122    type IntoIter = std::iter::Map<ChunkSeq<P>, F>;
123
124    fn into_iter(self) -> Self::IntoIter {
125        let chunks = ChunkSeq {
126            chunk_size: self.chunk_size,
127            len: self.len,
128            inner: if self.len > 0 { Some(self.base) } else { None },
129        };
130        chunks.map(self.map)
131    }
132
133    fn split_at(self, index: usize) -> (Self, Self) {
134        let elem_index = Ord::min(index * self.chunk_size, self.len);
135        let (left, right) = self.base.split_at(elem_index);
136        (
137            ChunkProducer {
138                chunk_size: self.chunk_size,
139                len: elem_index,
140                base: left,
141                map: self.map.clone(),
142            },
143            ChunkProducer {
144                chunk_size: self.chunk_size,
145                len: self.len - elem_index,
146                base: right,
147                map: self.map,
148            },
149        )
150    }
151
152    fn min_len(&self) -> usize {
153        div_round_up(self.base.min_len(), self.chunk_size)
154    }
155
156    fn max_len(&self) -> usize {
157        self.base.max_len() / self.chunk_size
158    }
159}
160
161pub(super) struct ChunkSeq<P> {
162    chunk_size: usize,
163    len: usize,
164    inner: Option<P>,
165}
166
167impl<P> Iterator for ChunkSeq<P>
168where
169    P: Producer,
170{
171    type Item = P::IntoIter;
172
173    fn next(&mut self) -> Option<Self::Item> {
174        let producer = self.inner.take()?;
175        if self.len > self.chunk_size {
176            let (left, right) = producer.split_at(self.chunk_size);
177            self.inner = Some(right);
178            self.len -= self.chunk_size;
179            Some(left.into_iter())
180        } else {
181            debug_assert!(self.len > 0);
182            self.len = 0;
183            Some(producer.into_iter())
184        }
185    }
186
187    fn size_hint(&self) -> (usize, Option<usize>) {
188        let len = self.len();
189        (len, Some(len))
190    }
191}
192
193impl<P> ExactSizeIterator for ChunkSeq<P>
194where
195    P: Producer,
196{
197    #[inline]
198    fn len(&self) -> usize {
199        div_round_up(self.len, self.chunk_size)
200    }
201}
202
203impl<P> DoubleEndedIterator for ChunkSeq<P>
204where
205    P: Producer,
206{
207    fn next_back(&mut self) -> Option<Self::Item> {
208        let producer = self.inner.take()?;
209        if self.len > self.chunk_size {
210            let mut size = self.len % self.chunk_size;
211            if size == 0 {
212                size = self.chunk_size;
213            }
214            let (left, right) = producer.split_at(self.len - size);
215            self.inner = Some(left);
216            self.len -= size;
217            Some(right.into_iter())
218        } else {
219            debug_assert!(self.len > 0);
220            self.len = 0;
221            Some(producer.into_iter())
222        }
223    }
224}
rayon/iter/chunks.rs

rayon/iter/
chunks.rs
