rayon/slice/
chunks.rs

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