rayon/
vec.rs

1//! Parallel iterator types for [vectors] (`Vec<T>`)
2//!
3//! You will rarely need to interact with this module directly unless you need
4//! to name one of the iterator types.
5//!
6//! [vectors]: mod@std::vec
7
8use crate::iter::plumbing::*;
9use crate::iter::*;
10use crate::math::simplify_range;
11use crate::slice::{Iter, IterMut};
12use std::iter;
13use std::mem;
14use std::ops::{Range, RangeBounds};
15use std::ptr;
16use std::slice;
17
18impl<'data, T: Sync + 'data> IntoParallelIterator for &'data Vec<T> {
19    type Item = &'data T;
20    type Iter = Iter<'data, T>;
21
22    fn into_par_iter(self) -> Self::Iter {
23        <&[T]>::into_par_iter(self)
24    }
25}
26
27impl<'data, T: Send + 'data> IntoParallelIterator for &'data mut Vec<T> {
28    type Item = &'data mut T;
29    type Iter = IterMut<'data, T>;
30
31    fn into_par_iter(self) -> Self::Iter {
32        <&mut [T]>::into_par_iter(self)
33    }
34}
35
36/// Parallel iterator that moves out of a vector.
37#[derive(Debug, Clone)]
38pub struct IntoIter<T> {
39    vec: Vec<T>,
40}
41
42impl<T: Send> IntoParallelIterator for Vec<T> {
43    type Item = T;
44    type Iter = IntoIter<T>;
45
46    fn into_par_iter(self) -> Self::Iter {
47        IntoIter { vec: self }
48    }
49}
50
51impl<T: Send> IntoParallelIterator for Box<[T]> {
52    type Item = T;
53    type Iter = IntoIter<T>;
54
55    fn into_par_iter(self) -> Self::Iter {
56        IntoIter { vec: self.into() }
57    }
58}
59
60impl<T: Send> ParallelIterator for IntoIter<T> {
61    type Item = T;
62
63    fn drive_unindexed<C>(self, consumer: C) -> C::Result
64    where
65        C: UnindexedConsumer<Self::Item>,
66    {
67        bridge(self, consumer)
68    }
69
70    fn opt_len(&self) -> Option<usize> {
71        Some(self.len())
72    }
73}
74
75impl<T: Send> IndexedParallelIterator for IntoIter<T> {
76    fn drive<C>(self, consumer: C) -> C::Result
77    where
78        C: Consumer<Self::Item>,
79    {
80        bridge(self, consumer)
81    }
82
83    fn len(&self) -> usize {
84        self.vec.len()
85    }
86
87    fn with_producer<CB>(mut self, callback: CB) -> CB::Output
88    where
89        CB: ProducerCallback<Self::Item>,
90    {
91        // Drain every item, and then the vector only needs to free its buffer.
92        self.vec.par_drain(..).with_producer(callback)
93    }
94}
95
96impl<'data, T: Send> ParallelDrainRange<usize> for &'data mut Vec<T> {
97    type Iter = Drain<'data, T>;
98    type Item = T;
99
100    fn par_drain<R: RangeBounds<usize>>(self, range: R) -> Self::Iter {
101        Drain {
102            orig_len: self.len(),
103            range: simplify_range(range, self.len()),
104            vec: self,
105        }
106    }
107}
108
109/// Draining parallel iterator that moves a range out of a vector, but keeps the total capacity.
110#[derive(Debug)]
111pub struct Drain<'data, T: Send> {
112    vec: &'data mut Vec<T>,
113    range: Range<usize>,
114    orig_len: usize,
115}
116
117impl<'data, T: Send> ParallelIterator for Drain<'data, T> {
118    type Item = T;
119
120    fn drive_unindexed<C>(self, consumer: C) -> C::Result
121    where
122        C: UnindexedConsumer<Self::Item>,
123    {
124        bridge(self, consumer)
125    }
126
127    fn opt_len(&self) -> Option<usize> {
128        Some(self.len())
129    }
130}
131
132impl<'data, T: Send> IndexedParallelIterator for Drain<'data, T> {
133    fn drive<C>(self, consumer: C) -> C::Result
134    where
135        C: Consumer<Self::Item>,
136    {
137        bridge(self, consumer)
138    }
139
140    fn len(&self) -> usize {
141        self.range.len()
142    }
143
144    fn with_producer<CB>(self, callback: CB) -> CB::Output
145    where
146        CB: ProducerCallback<Self::Item>,
147    {
148        unsafe {
149            // Make the vector forget about the drained items, and temporarily the tail too.
150            self.vec.set_len(self.range.start);
151
152            // Create the producer as the exclusive "owner" of the slice.
153            let producer = DrainProducer::from_vec(self.vec, self.range.len());
154
155            // The producer will move or drop each item from the drained range.
156            callback.callback(producer)
157        }
158    }
159}
160
161impl<'data, T: Send> Drop for Drain<'data, T> {
162    fn drop(&mut self) {
163        let Range { start, end } = self.range;
164        if self.vec.len() == self.orig_len {
165            // We must not have produced, so just call a normal drain to remove the items.
166            self.vec.drain(start..end);
167        } else if start == end {
168            // Empty range, so just restore the length to its original state
169            unsafe {
170                self.vec.set_len(self.orig_len);
171            }
172        } else if end < self.orig_len {
173            // The producer was responsible for consuming the drained items.
174            // Move the tail items to their new place, then set the length to include them.
175            unsafe {
176                let ptr = self.vec.as_mut_ptr().add(start);
177                let tail_ptr = self.vec.as_ptr().add(end);
178                let tail_len = self.orig_len - end;
179                ptr::copy(tail_ptr, ptr, tail_len);
180                self.vec.set_len(start + tail_len);
181            }
182        }
183    }
184}
185
186// ////////////////////////////////////////////////////////////////////////
187
188pub(crate) struct DrainProducer<'data, T: Send> {
189    slice: &'data mut [T],
190}
191
192impl<T: Send> DrainProducer<'_, T> {
193    /// Creates a draining producer, which *moves* items from the slice.
194    ///
195    /// Unsafe because `!Copy` data must not be read after the borrow is released.
196    pub(crate) unsafe fn new(slice: &mut [T]) -> DrainProducer<'_, T> {
197        DrainProducer { slice }
198    }
199
200    /// Creates a draining producer, which *moves* items from the tail of the vector.
201    ///
202    /// Unsafe because we're moving from beyond `vec.len()`, so the caller must ensure
203    /// that data is initialized and not read after the borrow is released.
204    unsafe fn from_vec(vec: &mut Vec<T>, len: usize) -> DrainProducer<'_, T> {
205        let start = vec.len();
206        assert!(vec.capacity() - start >= len);
207
208        // The pointer is derived from `Vec` directly, not through a `Deref`,
209        // so it has provenance over the whole allocation.
210        let ptr = vec.as_mut_ptr().add(start);
211        DrainProducer::new(slice::from_raw_parts_mut(ptr, len))
212    }
213}
214
215impl<'data, T: 'data + Send> Producer for DrainProducer<'data, T> {
216    type Item = T;
217    type IntoIter = SliceDrain<'data, T>;
218
219    fn into_iter(mut self) -> Self::IntoIter {
220        // replace the slice so we don't drop it twice
221        let slice = mem::take(&mut self.slice);
222        SliceDrain {
223            iter: slice.iter_mut(),
224        }
225    }
226
227    fn split_at(mut self, index: usize) -> (Self, Self) {
228        // replace the slice so we don't drop it twice
229        let slice = mem::take(&mut self.slice);
230        let (left, right) = slice.split_at_mut(index);
231        unsafe { (DrainProducer::new(left), DrainProducer::new(right)) }
232    }
233}
234
235impl<'data, T: 'data + Send> Drop for DrainProducer<'data, T> {
236    fn drop(&mut self) {
237        // extract the slice so we can use `Drop for [T]`
238        let slice_ptr: *mut [T] = mem::take::<&'data mut [T]>(&mut self.slice);
239        unsafe { ptr::drop_in_place::<[T]>(slice_ptr) };
240    }
241}
242
243// ////////////////////////////////////////////////////////////////////////
244
245// like std::vec::Drain, without updating a source Vec
246pub(crate) struct SliceDrain<'data, T> {
247    iter: slice::IterMut<'data, T>,
248}
249
250impl<'data, T: 'data> Iterator for SliceDrain<'data, T> {
251    type Item = T;
252
253    fn next(&mut self) -> Option<T> {
254        // Coerce the pointer early, so we don't keep the
255        // reference that's about to be invalidated.
256        let ptr: *const T = self.iter.next()?;
257        Some(unsafe { ptr::read(ptr) })
258    }
259
260    fn size_hint(&self) -> (usize, Option<usize>) {
261        self.iter.size_hint()
262    }
263
264    fn count(self) -> usize {
265        self.iter.len()
266    }
267}
268
269impl<'data, T: 'data> DoubleEndedIterator for SliceDrain<'data, T> {
270    fn next_back(&mut self) -> Option<Self::Item> {
271        // Coerce the pointer early, so we don't keep the
272        // reference that's about to be invalidated.
273        let ptr: *const T = self.iter.next_back()?;
274        Some(unsafe { ptr::read(ptr) })
275    }
276}
277
278impl<'data, T: 'data> ExactSizeIterator for SliceDrain<'data, T> {
279    fn len(&self) -> usize {
280        self.iter.len()
281    }
282}
283
284impl<'data, T: 'data> iter::FusedIterator for SliceDrain<'data, T> {}
285
286impl<'data, T: 'data> Drop for SliceDrain<'data, T> {
287    fn drop(&mut self) {
288        // extract the iterator so we can use `Drop for [T]`
289        let slice_ptr: *mut [T] = mem::replace(&mut self.iter, [].iter_mut()).into_slice();
290        unsafe { ptr::drop_in_place::<[T]>(slice_ptr) };
291    }
292}
rayon/vec.rs

rayon/
vec.rs
