rayon/slice/
chunks.rs

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

rayon/slice/
chunks.rs
