moirai_iter/parallel/
sources.rs

1use super::{
2    CollectConsumer, Consumer, IndexedParallelIterator, IntoParallelIterator,
3    IntoParallelRefIterator, ParallelExtend, ParallelIterator,
4};
5
6fn move_vec_items_into<T>(source: Vec<T>, target: &mut Vec<T>) {
7    target.clear();
8    let len = source.len();
9    if target.capacity() < len {
10        *target = source;
11        return;
12    }
13
14    // Consuming `source` moves every element without a `Clone` bound and
15    // releases its backing allocation while retaining `target`'s capacity.
16    // The prior `ManuallyDrop` copy leaked the source buffer.
17    target.extend(source);
18}
19
20/// Parallel iterator over a vector.
21pub struct VecParIter<T> {
22    data: Vec<T>,
23}
24
25impl<T> VecParIter<T> {
26    pub fn new(data: Vec<T>) -> Self {
27        Self { data }
28    }
29
30    pub(in crate::parallel) fn into_vec(self) -> Vec<T> {
31        self.data
32    }
33}
34
35impl<T: Send + Sync + 'static> ParallelIterator for VecParIter<T> {
36    type Item = T;
37
38    fn seq_items(self) -> Vec<Self::Item> {
39        self.data
40    }
41
42    fn drive<C, R>(mut self, consumer: C) -> R
43    where
44        C: Consumer<Self::Item, Result = R> + Send + Sync,
45        R: Send,
46    {
47        if self.data.len() <= 1 {
48            return consumer.consume(SequentialIterAdapter::new(self.data.into_iter()));
49        }
50
51        let mid = self.data.len() / 2;
52        let right_data = self.data.split_off(mid);
53        let left_data = std::mem::take(&mut self.data);
54
55        let (left_consumer, right_consumer) = consumer.split_at(left_data.len());
56
57        let left_result = left_consumer.consume(VecParIter::new(left_data));
58        let right_result = right_consumer.consume(VecParIter::new(right_data));
59
60        C::combine(left_result, right_result)
61    }
62}
63
64impl<T: Send + Sync + 'static> IndexedParallelIterator for VecParIter<T> {
65    fn len(&self) -> usize {
66        self.data.len()
67    }
68
69    fn collect_into_vec(self, target: &mut Vec<Self::Item>) {
70        move_vec_items_into(self.data, target);
71    }
72}
73
74/// Range parallel iterator.
75pub struct RangeParIter<T> {
76    start: T,
77    end: T,
78}
79
80impl<T> RangeParIter<T>
81where
82    T: Send + Sync + Clone + 'static + PartialOrd + std::ops::Add<Output = T> + From<u8>,
83{
84    pub fn new(start: T, end: T) -> Self {
85        Self { start, end }
86    }
87}
88
89impl<T> ParallelIterator for RangeParIter<T>
90where
91    T: Send + Sync + Clone + 'static + PartialOrd + std::ops::Add<Output = T> + From<u8>,
92{
93    type Item = T;
94
95    fn seq_items(self) -> Vec<Self::Item> {
96        let mut items = Vec::new();
97        let mut current = self.start;
98        while current < self.end {
99            items.push(current.clone());
100            current = current + T::from(1u8);
101        }
102        items
103    }
104
105    fn drive<C, R>(self, consumer: C) -> R
106    where
107        C: Consumer<Self::Item, Result = R> + Send + Sync,
108        R: Send,
109    {
110        let mut items = Vec::new();
111        let mut current = self.start;
112        while current < self.end {
113            items.push(current.clone());
114            current = current + T::from(1u8);
115        }
116
117        VecParIter::new(items).drive(consumer)
118    }
119}
120
121impl IndexedParallelIterator for RangeParIter<usize> {
122    fn len(&self) -> usize {
123        self.end.saturating_sub(self.start)
124    }
125
126    fn collect_into_vec(self, target: &mut Vec<Self::Item>) {
127        target.clear();
128        target.extend(self.start..self.end);
129    }
130}
131
132/// Sequential iterator adapter for compatibility.
133pub struct SequentialAdapter<I> {
134    iter: I,
135}
136
137impl<I> SequentialAdapter<I> {
138    pub(super) fn new(iter: I) -> Self {
139        Self { iter }
140    }
141}
142
143pub struct SequentialIterAdapter<I> {
144    iter: I,
145}
146
147impl<I> SequentialIterAdapter<I> {
148    pub(super) fn new(iter: I) -> Self {
149        Self { iter }
150    }
151}
152
153impl<I> ParallelIterator for SequentialIterAdapter<I>
154where
155    I: Iterator + Send,
156    I::Item: Send + Sync + 'static,
157{
158    type Item = I::Item;
159
160    fn seq_items(self) -> Vec<Self::Item> {
161        self.iter.collect()
162    }
163
164    fn drive<C, R>(self, consumer: C) -> R
165    where
166        C: Consumer<Self::Item, Result = R> + Send + Sync,
167        R: Send,
168    {
169        let items: Vec<Self::Item> = self.iter.collect();
170        consumer.consume(VecParIter::new(items))
171    }
172}
173
174impl<I> IndexedParallelIterator for SequentialIterAdapter<I>
175where
176    I: ExactSizeIterator + Send,
177    I::Item: Send + Sync + 'static,
178{
179    fn len(&self) -> usize {
180        self.iter.len()
181    }
182
183    fn collect_into_vec(self, target: &mut Vec<Self::Item>) {
184        target.clear();
185        target.extend(self.iter);
186    }
187}
188
189impl<T: Send + Sync + 'static> IntoParallelIterator for Vec<T> {
190    type Item = T;
191    type Iter = VecParIter<T>;
192
193    fn into_par_iter(self) -> Self::Iter {
194        VecParIter::new(self)
195    }
196}
197
198impl<'data, T: Send + Sync + 'data> IntoParallelRefIterator<'data> for Vec<T> {
199    type Item = &'data T;
200    type Iter = VecRefParIter<'data, T>;
201
202    fn par_iter(&'data self) -> Self::Iter {
203        VecRefParIter::new(self)
204    }
205}
206
207/// Parallel iterator over vector references.
208pub struct VecRefParIter<'data, T> {
209    data: &'data Vec<T>,
210}
211
212impl<'data, T> VecRefParIter<'data, T> {
213    fn new(data: &'data Vec<T>) -> Self {
214        Self { data }
215    }
216
217    pub(in crate::parallel) fn into_slice(self) -> &'data [T] {
218        self.data.as_slice()
219    }
220
221    /// Return matching logical positions without materializing borrowed items.
222    pub fn positions<F>(self, predicate: F) -> VecRefPositions<'data, T, F>
223    where
224        F: Fn(&'data T) -> bool + Send + Sync + Clone,
225    {
226        VecRefPositions {
227            data: self.data,
228            predicate,
229        }
230    }
231}
232
233impl<'data, T: Send + Sync + 'data> ParallelIterator for VecRefParIter<'data, T> {
234    type Item = &'data T;
235
236    fn seq_items(self) -> Vec<Self::Item> {
237        self.data.iter().collect()
238    }
239
240    fn drive<C, R>(self, consumer: C) -> R
241    where
242        C: Consumer<Self::Item, Result = R> + Send + Sync,
243        R: Send,
244    {
245        let refs: Vec<&'data T> = self.data.iter().collect();
246        consumer.consume(RefVecParIter::new(refs))
247    }
248}
249
250impl<'data, T: Send + Sync + 'data> IndexedParallelIterator for VecRefParIter<'data, T> {
251    fn len(&self) -> usize {
252        self.data.len()
253    }
254
255    fn collect_into_vec(self, target: &mut Vec<Self::Item>) {
256        target.clear();
257        target.extend(self.data.iter());
258    }
259}
260
261/// Position stream over borrowed vector storage.
262pub struct VecRefPositions<'data, T, F> {
263    data: &'data Vec<T>,
264    predicate: F,
265}
266
267impl<'data, T, F> ParallelIterator for VecRefPositions<'data, T, F>
268where
269    T: Send + Sync + 'data,
270    F: Fn(&'data T) -> bool + Send + Sync + Clone,
271{
272    type Item = usize;
273
274    fn seq_items(self) -> Vec<Self::Item> {
275        self.data
276            .iter()
277            .enumerate()
278            .filter_map(|(index, item)| (self.predicate)(item).then_some(index))
279            .collect()
280    }
281
282    fn drive<C, R>(self, consumer: C) -> R
283    where
284        C: Consumer<Self::Item, Result = R> + Send + Sync,
285        R: Send,
286    {
287        consumer.consume(VecParIter::new(self.seq_items()))
288    }
289}
290
291/// A parallel iterator specifically for reference vectors.
292pub struct RefVecParIter<'a, T> {
293    data: Vec<&'a T>,
294}
295
296impl<'a, T> RefVecParIter<'a, T> {
297    fn new(data: Vec<&'a T>) -> Self {
298        Self { data }
299    }
300}
301
302impl<'a, T: Send + Sync> ParallelIterator for RefVecParIter<'a, T> {
303    type Item = &'a T;
304
305    fn seq_items(self) -> Vec<Self::Item> {
306        self.data
307    }
308
309    fn drive<C, R>(mut self, consumer: C) -> R
310    where
311        C: Consumer<Self::Item, Result = R> + Send + Sync,
312        R: Send,
313    {
314        // Sequential base case: consumer.consume(RefVecParIter) is safe because
315        // CollectConsumer::consume now calls seq_items(), terminating the chain.
316        if self.data.len() <= 1 {
317            return consumer.consume(RefVecParIter::new(self.data));
318        }
319
320        let mid = self.data.len() / 2;
321        let right_data = self.data.split_off(mid);
322        let left_data = std::mem::take(&mut self.data);
323
324        let (left_consumer, right_consumer) = consumer.split_at(left_data.len());
325
326        let left_result = left_consumer.consume(RefVecParIter::new(left_data));
327        let right_result = right_consumer.consume(RefVecParIter::new(right_data));
328
329        C::combine(left_result, right_result)
330    }
331}
332
333impl<'a, T: Send + Sync> IndexedParallelIterator for RefVecParIter<'a, T> {
334    fn len(&self) -> usize {
335        self.data.len()
336    }
337
338    fn collect_into_vec(self, target: &mut Vec<Self::Item>) {
339        move_vec_items_into(self.data, target);
340    }
341}
342
343impl IntoParallelIterator for std::ops::Range<usize> {
344    type Item = usize;
345    type Iter = RangeParIter<usize>;
346
347    fn into_par_iter(self) -> Self::Iter {
348        RangeParIter::new(self.start, self.end)
349    }
350}
351
352impl<T: Send + Sync> ParallelExtend<T> for Vec<T> {
353    fn par_extend<I>(&mut self, par_iter: I)
354    where
355        I: ParallelIterator<Item = T>,
356    {
357        // Drive through CollectConsumer rather than calling collect::<Vec<_>>(), which
358        // would call par_extend again and create infinite mutual recursion.
359        self.extend(par_iter.drive(CollectConsumer::new()));
360    }
361}
moirai_iter/parallel/sources.rs

moirai_iter/parallel/
sources.rs
