rayon/slice/
chunks.rs

1use crate::iter::plumbing::*;
2use crate::iter::*;
3
4/// Parallel iterator over immutable non-overlapping chunks of a slice
5#[derive(Debug)]
6pub struct Chunks<'data, T> {
7    chunk_size: usize,
8    slice: &'data [T],
9}
10
11impl<'data, T> Chunks<'data, T> {
12    pub(super) fn new(chunk_size: usize, slice: &'data [T]) -> Self {
13        Self { chunk_size, slice }
14    }
15}
16
17impl<T> Clone for Chunks<'_, T> {
18    fn clone(&self) -> Self {
19        Chunks { ..*self }
20    }
21}
22
23impl<'data, T: Sync> ParallelIterator for Chunks<'data, T> {
24    type Item = &'data [T];
25
26    fn drive_unindexed<C>(self, consumer: C) -> C::Result
27    where
28        C: UnindexedConsumer<Self::Item>,
29    {
30        bridge(self, consumer)
31    }
32
33    fn opt_len(&self) -> Option<usize> {
34        Some(self.len())
35    }
36}
37
38impl<T: Sync> IndexedParallelIterator for Chunks<'_, T> {
39    fn drive<C>(self, consumer: C) -> C::Result
40    where
41        C: Consumer<Self::Item>,
42    {
43        bridge(self, consumer)
44    }
45
46    fn len(&self) -> usize {
47        self.slice.len().div_ceil(self.chunk_size)
48    }
49
50    fn with_producer<CB>(self, callback: CB) -> CB::Output
51    where
52        CB: ProducerCallback<Self::Item>,
53    {
54        callback.callback(ChunksProducer {
55            chunk_size: self.chunk_size,
56            slice: self.slice,
57        })
58    }
59}
60
61struct ChunksProducer<'data, T: Sync> {
62    chunk_size: usize,
63    slice: &'data [T],
64}
65
66impl<'data, T: 'data + Sync> Producer for ChunksProducer<'data, T> {
67    type Item = &'data [T];
68    type IntoIter = ::std::slice::Chunks<'data, T>;
69
70    fn into_iter(self) -> Self::IntoIter {
71        self.slice.chunks(self.chunk_size)
72    }
73
74    fn split_at(self, index: usize) -> (Self, Self) {
75        let elem_index = Ord::min(index * self.chunk_size, self.slice.len());
76        let (left, right) = self.slice.split_at(elem_index);
77        (
78            ChunksProducer {
79                chunk_size: self.chunk_size,
80                slice: left,
81            },
82            ChunksProducer {
83                chunk_size: self.chunk_size,
84                slice: right,
85            },
86        )
87    }
88}
89
90/// Parallel iterator over immutable non-overlapping chunks of a slice
91#[derive(Debug)]
92pub struct ChunksExact<'data, T> {
93    chunk_size: usize,
94    slice: &'data [T],
95    rem: &'data [T],
96}
97
98impl<'data, T> ChunksExact<'data, T> {
99    pub(super) fn new(chunk_size: usize, slice: &'data [T]) -> Self {
100        let rem_len = slice.len() % chunk_size;
101        let len = slice.len() - rem_len;
102        let (slice, rem) = slice.split_at(len);
103        Self {
104            chunk_size,
105            slice,
106            rem,
107        }
108    }
109
110    /// Return the remainder of the original slice that is not going to be
111    /// returned by the iterator. The returned slice has at most `chunk_size-1`
112    /// elements.
113    pub fn remainder(&self) -> &'data [T] {
114        self.rem
115    }
116}
117
118impl<T> Clone for ChunksExact<'_, T> {
119    fn clone(&self) -> Self {
120        ChunksExact { ..*self }
121    }
122}
123
124impl<'data, T: Sync> ParallelIterator for ChunksExact<'data, T> {
125    type Item = &'data [T];
126
127    fn drive_unindexed<C>(self, consumer: C) -> C::Result
128    where
129        C: UnindexedConsumer<Self::Item>,
130    {
131        bridge(self, consumer)
132    }
133
134    fn opt_len(&self) -> Option<usize> {
135        Some(self.len())
136    }
137}
138
139impl<T: Sync> IndexedParallelIterator for ChunksExact<'_, T> {
140    fn drive<C>(self, consumer: C) -> C::Result
141    where
142        C: Consumer<Self::Item>,
143    {
144        bridge(self, consumer)
145    }
146
147    fn len(&self) -> usize {
148        self.slice.len() / self.chunk_size
149    }
150
151    fn with_producer<CB>(self, callback: CB) -> CB::Output
152    where
153        CB: ProducerCallback<Self::Item>,
154    {
155        callback.callback(ChunksExactProducer {
156            chunk_size: self.chunk_size,
157            slice: self.slice,
158        })
159    }
160}
161
162struct ChunksExactProducer<'data, T: Sync> {
163    chunk_size: usize,
164    slice: &'data [T],
165}
166
167impl<'data, T: 'data + Sync> Producer for ChunksExactProducer<'data, T> {
168    type Item = &'data [T];
169    type IntoIter = ::std::slice::ChunksExact<'data, T>;
170
171    fn into_iter(self) -> Self::IntoIter {
172        self.slice.chunks_exact(self.chunk_size)
173    }
174
175    fn split_at(self, index: usize) -> (Self, Self) {
176        let elem_index = index * self.chunk_size;
177        let (left, right) = self.slice.split_at(elem_index);
178        (
179            ChunksExactProducer {
180                chunk_size: self.chunk_size,
181                slice: left,
182            },
183            ChunksExactProducer {
184                chunk_size: self.chunk_size,
185                slice: right,
186            },
187        )
188    }
189}
190
191/// Parallel iterator over mutable non-overlapping chunks of a slice
192#[derive(Debug)]
193pub struct ChunksMut<'data, T> {
194    chunk_size: usize,
195    slice: &'data mut [T],
196}
197
198impl<'data, T> ChunksMut<'data, T> {
199    pub(super) fn new(chunk_size: usize, slice: &'data mut [T]) -> Self {
200        Self { chunk_size, slice }
201    }
202}
203
204impl<'data, T: Send> ParallelIterator for ChunksMut<'data, T> {
205    type Item = &'data mut [T];
206
207    fn drive_unindexed<C>(self, consumer: C) -> C::Result
208    where
209        C: UnindexedConsumer<Self::Item>,
210    {
211        bridge(self, consumer)
212    }
213
214    fn opt_len(&self) -> Option<usize> {
215        Some(self.len())
216    }
217}
218
219impl<T: Send> IndexedParallelIterator for ChunksMut<'_, T> {
220    fn drive<C>(self, consumer: C) -> C::Result
221    where
222        C: Consumer<Self::Item>,
223    {
224        bridge(self, consumer)
225    }
226
227    fn len(&self) -> usize {
228        self.slice.len().div_ceil(self.chunk_size)
229    }
230
231    fn with_producer<CB>(self, callback: CB) -> CB::Output
232    where
233        CB: ProducerCallback<Self::Item>,
234    {
235        callback.callback(ChunksMutProducer {
236            chunk_size: self.chunk_size,
237            slice: self.slice,
238        })
239    }
240}
241
242struct ChunksMutProducer<'data, T: Send> {
243    chunk_size: usize,
244    slice: &'data mut [T],
245}
246
247impl<'data, T: 'data + Send> Producer for ChunksMutProducer<'data, T> {
248    type Item = &'data mut [T];
249    type IntoIter = ::std::slice::ChunksMut<'data, T>;
250
251    fn into_iter(self) -> Self::IntoIter {
252        self.slice.chunks_mut(self.chunk_size)
253    }
254
255    fn split_at(self, index: usize) -> (Self, Self) {
256        let elem_index = Ord::min(index * self.chunk_size, self.slice.len());
257        let (left, right) = self.slice.split_at_mut(elem_index);
258        (
259            ChunksMutProducer {
260                chunk_size: self.chunk_size,
261                slice: left,
262            },
263            ChunksMutProducer {
264                chunk_size: self.chunk_size,
265                slice: right,
266            },
267        )
268    }
269}
270
271/// Parallel iterator over mutable non-overlapping chunks of a slice
272#[derive(Debug)]
273pub struct ChunksExactMut<'data, T> {
274    chunk_size: usize,
275    slice: &'data mut [T],
276    rem: &'data mut [T],
277}
278
279impl<'data, T> ChunksExactMut<'data, T> {
280    pub(super) fn new(chunk_size: usize, slice: &'data mut [T]) -> Self {
281        let rem_len = slice.len() % chunk_size;
282        let len = slice.len() - rem_len;
283        let (slice, rem) = slice.split_at_mut(len);
284        Self {
285            chunk_size,
286            slice,
287            rem,
288        }
289    }
290
291    /// Return the remainder of the original slice that is not going to be
292    /// returned by the iterator. The returned slice has at most `chunk_size-1`
293    /// elements.
294    ///
295    /// Note that this has to consume `self` to return the original lifetime of
296    /// the data, which prevents this from actually being used as a parallel
297    /// iterator since that also consumes. This method is provided for parity
298    /// with `std::iter::ChunksExactMut`, but consider calling `remainder()` or
299    /// `take_remainder()` as alternatives.
300    pub fn into_remainder(self) -> &'data mut [T] {
301        self.rem
302    }
303
304    /// Return the remainder of the original slice that is not going to be
305    /// returned by the iterator. The returned slice has at most `chunk_size-1`
306    /// elements.
307    ///
308    /// Consider `take_remainder()` if you need access to the data with its
309    /// original lifetime, rather than borrowing through `&mut self` here.
310    pub fn remainder(&mut self) -> &mut [T] {
311        self.rem
312    }
313
314    /// Return the remainder of the original slice that is not going to be
315    /// returned by the iterator. The returned slice has at most `chunk_size-1`
316    /// elements. Subsequent calls will return an empty slice.
317    pub fn take_remainder(&mut self) -> &'data mut [T] {
318        std::mem::take(&mut self.rem)
319    }
320}
321
322impl<'data, T: Send> ParallelIterator for ChunksExactMut<'data, T> {
323    type Item = &'data mut [T];
324
325    fn drive_unindexed<C>(self, consumer: C) -> C::Result
326    where
327        C: UnindexedConsumer<Self::Item>,
328    {
329        bridge(self, consumer)
330    }
331
332    fn opt_len(&self) -> Option<usize> {
333        Some(self.len())
334    }
335}
336
337impl<T: Send> IndexedParallelIterator for ChunksExactMut<'_, T> {
338    fn drive<C>(self, consumer: C) -> C::Result
339    where
340        C: Consumer<Self::Item>,
341    {
342        bridge(self, consumer)
343    }
344
345    fn len(&self) -> usize {
346        self.slice.len() / self.chunk_size
347    }
348
349    fn with_producer<CB>(self, callback: CB) -> CB::Output
350    where
351        CB: ProducerCallback<Self::Item>,
352    {
353        callback.callback(ChunksExactMutProducer {
354            chunk_size: self.chunk_size,
355            slice: self.slice,
356        })
357    }
358}
359
360struct ChunksExactMutProducer<'data, T: Send> {
361    chunk_size: usize,
362    slice: &'data mut [T],
363}
364
365impl<'data, T: 'data + Send> Producer for ChunksExactMutProducer<'data, T> {
366    type Item = &'data mut [T];
367    type IntoIter = ::std::slice::ChunksExactMut<'data, T>;
368
369    fn into_iter(self) -> Self::IntoIter {
370        self.slice.chunks_exact_mut(self.chunk_size)
371    }
372
373    fn split_at(self, index: usize) -> (Self, Self) {
374        let elem_index = index * self.chunk_size;
375        let (left, right) = self.slice.split_at_mut(elem_index);
376        (
377            ChunksExactMutProducer {
378                chunk_size: self.chunk_size,
379                slice: left,
380            },
381            ChunksExactMutProducer {
382                chunk_size: self.chunk_size,
383                slice: right,
384            },
385        )
386    }
387}