rayon/slice/
chunk_by.rs

1use crate::iter::plumbing::*;
2use crate::iter::*;
3use std::fmt;
4use std::marker::PhantomData;
5
6trait ChunkBySlice<T>: AsRef<[T]> + Default + Send {
7    fn split(self, index: usize) -> (Self, Self);
8    fn chunk_by(self, pred: &impl Fn(&T, &T) -> bool) -> impl Iterator<Item = Self>;
9
10    fn find(&self, pred: &impl Fn(&T, &T) -> bool, start: usize, end: usize) -> Option<usize> {
11        self.as_ref()[start..end]
12            .windows(2)
13            .position(move |w| !pred(&w[0], &w[1]))
14            .map(|i| i + 1)
15    }
16
17    fn rfind(&self, pred: &impl Fn(&T, &T) -> bool, end: usize) -> Option<usize> {
18        self.as_ref()[..end]
19            .windows(2)
20            .rposition(move |w| !pred(&w[0], &w[1]))
21            .map(|i| i + 1)
22    }
23}
24
25impl<T: Sync> ChunkBySlice<T> for &[T] {
26    fn split(self, index: usize) -> (Self, Self) {
27        self.split_at(index)
28    }
29
30    fn chunk_by(self, pred: &impl Fn(&T, &T) -> bool) -> impl Iterator<Item = Self> {
31        self.chunk_by(pred)
32    }
33}
34
35impl<T: Send> ChunkBySlice<T> for &mut [T] {
36    fn split(self, index: usize) -> (Self, Self) {
37        self.split_at_mut(index)
38    }
39
40    fn chunk_by(self, pred: &impl Fn(&T, &T) -> bool) -> impl Iterator<Item = Self> {
41        self.chunk_by_mut(pred)
42    }
43}
44
45struct ChunkByProducer<'p, T, Slice, Pred> {
46    slice: Slice,
47    pred: &'p Pred,
48    tail: usize,
49    marker: PhantomData<fn(&T)>,
50}
51
52// Note: this implementation is very similar to `SplitProducer`.
53impl<T, Slice, Pred> UnindexedProducer for ChunkByProducer<'_, T, Slice, Pred>
54where
55    Slice: ChunkBySlice<T>,
56    Pred: Fn(&T, &T) -> bool + Send + Sync,
57{
58    type Item = Slice;
59
60    fn split(self) -> (Self, Option<Self>) {
61        if self.tail < 2 {
62            return (Self { tail: 0, ..self }, None);
63        }
64
65        // Look forward for the separator, and failing that look backward.
66        let mid = self.tail / 2;
67        let index = match self.slice.find(self.pred, mid, self.tail) {
68            Some(i) => Some(mid + i),
69            None => self.slice.rfind(self.pred, mid + 1),
70        };
71
72        if let Some(index) = index {
73            let (left, right) = self.slice.split(index);
74
75            let (left_tail, right_tail) = if index <= mid {
76                // If we scanned backwards to find the separator, everything in
77                // the right side is exhausted, with no separators left to find.
78                (index, 0)
79            } else {
80                (mid + 1, self.tail - index)
81            };
82
83            // Create the left split before the separator.
84            let left = Self {
85                slice: left,
86                tail: left_tail,
87                ..self
88            };
89
90            // Create the right split following the separator.
91            let right = Self {
92                slice: right,
93                tail: right_tail,
94                ..self
95            };
96
97            (left, Some(right))
98        } else {
99            // The search is exhausted, no more separators...
100            (Self { tail: 0, ..self }, None)
101        }
102    }
103
104    fn fold_with<F>(self, mut folder: F) -> F
105    where
106        F: Folder<Self::Item>,
107    {
108        let Self {
109            slice, pred, tail, ..
110        } = self;
111
112        let (slice, tail) = if tail == slice.as_ref().len() {
113            // No tail section, so just let `consume_iter` do it all.
114            (Some(slice), None)
115        } else if let Some(index) = slice.rfind(pred, tail) {
116            // We found the last separator to complete the tail, so
117            // end with that slice after `consume_iter` finds the rest.
118            let (left, right) = slice.split(index);
119            (Some(left), Some(right))
120        } else {
121            // We know there are no separators at all, so it's all "tail".
122            (None, Some(slice))
123        };
124
125        if let Some(slice) = slice {
126            folder = folder.consume_iter(slice.chunk_by(pred));
127        }
128
129        if let Some(tail) = tail {
130            folder = folder.consume(tail);
131        }
132
133        folder
134    }
135}
136
137/// Parallel iterator over slice in (non-overlapping) chunks separated by a predicate.
138///
139/// This struct is created by the [`par_chunk_by`] method on `&[T]`.
140///
141/// [`par_chunk_by`]: super::ParallelSlice::par_chunk_by()
142pub struct ChunkBy<'data, T, P> {
143    pred: P,
144    slice: &'data [T],
145}
146
147impl<T, P: Clone> Clone for ChunkBy<'_, T, P> {
148    fn clone(&self) -> Self {
149        ChunkBy {
150            pred: self.pred.clone(),
151            slice: self.slice,
152        }
153    }
154}
155
156impl<T: fmt::Debug, P> fmt::Debug for ChunkBy<'_, T, P> {
157    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
158        f.debug_struct("ChunkBy")
159            .field("slice", &self.slice)
160            .finish()
161    }
162}
163
164impl<'data, T, P> ChunkBy<'data, T, P> {
165    pub(super) fn new(slice: &'data [T], pred: P) -> Self {
166        Self { pred, slice }
167    }
168}
169
170impl<'data, T, P> ParallelIterator for ChunkBy<'data, T, P>
171where
172    T: Sync,
173    P: Fn(&T, &T) -> bool + Send + Sync,
174{
175    type Item = &'data [T];
176
177    fn drive_unindexed<C>(self, consumer: C) -> C::Result
178    where
179        C: UnindexedConsumer<Self::Item>,
180    {
181        bridge_unindexed(
182            ChunkByProducer {
183                tail: self.slice.len(),
184                slice: self.slice,
185                pred: &self.pred,
186                marker: PhantomData,
187            },
188            consumer,
189        )
190    }
191}
192
193/// Parallel iterator over slice in (non-overlapping) mutable chunks
194/// separated by a predicate.
195///
196/// This struct is created by the [`par_chunk_by_mut`] method on `&mut [T]`.
197///
198/// [`par_chunk_by_mut`]: super::ParallelSliceMut::par_chunk_by_mut()
199pub struct ChunkByMut<'data, T, P> {
200    pred: P,
201    slice: &'data mut [T],
202}
203
204impl<T: fmt::Debug, P> fmt::Debug for ChunkByMut<'_, T, P> {
205    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
206        f.debug_struct("ChunkByMut")
207            .field("slice", &self.slice)
208            .finish()
209    }
210}
211
212impl<'data, T, P> ChunkByMut<'data, T, P> {
213    pub(super) fn new(slice: &'data mut [T], pred: P) -> Self {
214        Self { pred, slice }
215    }
216}
217
218impl<'data, T, P> ParallelIterator for ChunkByMut<'data, T, P>
219where
220    T: Send,
221    P: Fn(&T, &T) -> bool + Send + Sync,
222{
223    type Item = &'data mut [T];
224
225    fn drive_unindexed<C>(self, consumer: C) -> C::Result
226    where
227        C: UnindexedConsumer<Self::Item>,
228    {
229        bridge_unindexed(
230            ChunkByProducer {
231                tail: self.slice.len(),
232                slice: self.slice,
233                pred: &self.pred,
234                marker: PhantomData,
235            },
236            consumer,
237        )
238    }
239}
rayon/slice/chunk_by.rs

rayon/slice/
chunk_by.rs
