rayon/slice/
chunk_by.rs

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

rayon/slice/
chunk_by.rs
