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