]>
Commit | Line | Data |
---|---|---|
2c00a5a8 XL |
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 | ||
f035d41b XL |
8 | use crate::iter::plumbing::*; |
9 | use crate::iter::*; | |
5869c6ff | 10 | use crate::math::simplify_range; |
17df50a5 | 11 | use crate::slice::{Iter, IterMut}; |
5869c6ff XL |
12 | use std::iter; |
13 | use std::mem; | |
14 | use std::ops::{Range, RangeBounds}; | |
15 | use std::ptr; | |
16 | use std::slice; | |
2c00a5a8 | 17 | |
17df50a5 XL |
18 | impl<'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 | ||
27 | impl<'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 | ||
2c00a5a8 XL |
36 | /// Parallel iterator that moves out of a vector. |
37 | #[derive(Debug, Clone)] | |
38 | pub struct IntoIter<T: Send> { | |
39 | vec: Vec<T>, | |
40 | } | |
41 | ||
42 | impl<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 | ||
51 | impl<T: Send> ParallelIterator for IntoIter<T> { | |
52 | type Item = T; | |
53 | ||
54 | fn drive_unindexed<C>(self, consumer: C) -> C::Result | |
416331ca XL |
55 | where |
56 | C: UnindexedConsumer<Self::Item>, | |
2c00a5a8 XL |
57 | { |
58 | bridge(self, consumer) | |
59 | } | |
60 | ||
61 | fn opt_len(&self) -> Option<usize> { | |
62 | Some(self.len()) | |
63 | } | |
64 | } | |
65 | ||
66 | impl<T: Send> IndexedParallelIterator for IntoIter<T> { | |
67 | fn drive<C>(self, consumer: C) -> C::Result | |
416331ca XL |
68 | where |
69 | C: Consumer<Self::Item>, | |
2c00a5a8 XL |
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 | |
416331ca XL |
79 | where |
80 | CB: ProducerCallback<Self::Item>, | |
2c00a5a8 | 81 | { |
5869c6ff XL |
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 | ||
87 | impl<'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)] | |
102 | pub struct Drain<'data, T: Send> { | |
103 | vec: &'data mut Vec<T>, | |
104 | range: Range<usize>, | |
105 | orig_len: usize, | |
106 | } | |
107 | ||
108 | impl<'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 | ||
123 | impl<'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 | ||
487cf647 | 135 | fn with_producer<CB>(self, callback: CB) -> CB::Output |
5869c6ff XL |
136 | where |
137 | CB: ProducerCallback<Self::Item>, | |
138 | { | |
2c00a5a8 | 139 | unsafe { |
5869c6ff | 140 | // Make the vector forget about the drained items, and temporarily the tail too. |
04454e1e | 141 | self.vec.set_len(self.range.start); |
2c00a5a8 | 142 | |
17df50a5 | 143 | // Create the producer as the exclusive "owner" of the slice. |
487cf647 | 144 | let producer = DrainProducer::from_vec(self.vec, self.range.len()); |
2c00a5a8 | 145 | |
5869c6ff | 146 | // The producer will move or drop each item from the drained range. |
17df50a5 | 147 | callback.callback(producer) |
5869c6ff XL |
148 | } |
149 | } | |
150 | } | |
151 | ||
152 | impl<'data, T: Send> Drop for Drain<'data, T> { | |
153 | fn drop(&mut self) { | |
487cf647 FG |
154 | let Range { start, end } = 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(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 = self.vec.as_mut_ptr().add(start); | |
168 | let tail_ptr = self.vec.as_ptr().add(end); | |
169 | let tail_len = self.orig_len - end; | |
170 | ptr::copy(tail_ptr, ptr, tail_len); | |
171 | self.vec.set_len(start + tail_len); | |
5869c6ff | 172 | } |
2c00a5a8 XL |
173 | } |
174 | } | |
175 | } | |
176 | ||
177 | /// //////////////////////////////////////////////////////////////////////// | |
178 | ||
5869c6ff | 179 | pub(crate) struct DrainProducer<'data, T: Send> { |
2c00a5a8 XL |
180 | slice: &'data mut [T], |
181 | } | |
182 | ||
04454e1e | 183 | impl<T: Send> DrainProducer<'_, T> { |
5869c6ff XL |
184 | /// Creates a draining producer, which *moves* items from the slice. |
185 | /// | |
487cf647 | 186 | /// Unsafe because `!Copy` data must not be read after the borrow is released. |
04454e1e | 187 | pub(crate) unsafe fn new(slice: &mut [T]) -> DrainProducer<'_, T> { |
5869c6ff XL |
188 | DrainProducer { slice } |
189 | } | |
04454e1e FG |
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 = 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 = vec.as_mut_ptr().add(start); | |
202 | DrainProducer::new(slice::from_raw_parts_mut(ptr, len)) | |
203 | } | |
5869c6ff XL |
204 | } |
205 | ||
206 | impl<'data, T: 'data + Send> Producer for DrainProducer<'data, T> { | |
2c00a5a8 XL |
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 | |
487cf647 | 212 | let slice = mem::take(&mut self.slice); |
416331ca XL |
213 | SliceDrain { |
214 | iter: slice.iter_mut(), | |
215 | } | |
2c00a5a8 XL |
216 | } |
217 | ||
218 | fn split_at(mut self, index: usize) -> (Self, Self) { | |
219 | // replace the slice so we don't drop it twice | |
487cf647 | 220 | let slice = mem::take(&mut self.slice); |
2c00a5a8 | 221 | let (left, right) = slice.split_at_mut(index); |
5869c6ff | 222 | unsafe { (DrainProducer::new(left), DrainProducer::new(right)) } |
2c00a5a8 XL |
223 | } |
224 | } | |
225 | ||
5869c6ff | 226 | impl<'data, T: 'data + Send> Drop for DrainProducer<'data, T> { |
2c00a5a8 | 227 | fn drop(&mut self) { |
5869c6ff XL |
228 | // use `Drop for [T]` |
229 | unsafe { ptr::drop_in_place(self.slice) }; | |
2c00a5a8 XL |
230 | } |
231 | } | |
232 | ||
233 | /// //////////////////////////////////////////////////////////////////////// | |
234 | ||
235 | // like std::vec::Drain, without updating a source Vec | |
5869c6ff XL |
236 | pub(crate) struct SliceDrain<'data, T> { |
237 | iter: slice::IterMut<'data, T>, | |
2c00a5a8 XL |
238 | } |
239 | ||
240 | impl<'data, T: 'data> Iterator for SliceDrain<'data, T> { | |
241 | type Item = T; | |
242 | ||
243 | fn next(&mut self) -> Option<T> { | |
17df50a5 XL |
244 | // Coerce the pointer early, so we don't keep the |
245 | // reference that's about to be invalidated. | |
246 | let ptr: *const T = self.iter.next()?; | |
5869c6ff | 247 | Some(unsafe { ptr::read(ptr) }) |
2c00a5a8 XL |
248 | } |
249 | ||
250 | fn size_hint(&self) -> (usize, Option<usize>) { | |
5869c6ff XL |
251 | self.iter.size_hint() |
252 | } | |
253 | ||
254 | fn count(self) -> usize { | |
255 | self.iter.len() | |
2c00a5a8 XL |
256 | } |
257 | } | |
258 | ||
259 | impl<'data, T: 'data> DoubleEndedIterator for SliceDrain<'data, T> { | |
260 | fn next_back(&mut self) -> Option<Self::Item> { | |
17df50a5 XL |
261 | // Coerce the pointer early, so we don't keep the |
262 | // reference that's about to be invalidated. | |
263 | let ptr: *const T = self.iter.next_back()?; | |
5869c6ff | 264 | Some(unsafe { ptr::read(ptr) }) |
2c00a5a8 XL |
265 | } |
266 | } | |
267 | ||
268 | impl<'data, T: 'data> ExactSizeIterator for SliceDrain<'data, T> { | |
269 | fn len(&self) -> usize { | |
270 | self.iter.len() | |
271 | } | |
272 | } | |
273 | ||
5869c6ff XL |
274 | impl<'data, T: 'data> iter::FusedIterator for SliceDrain<'data, T> {} |
275 | ||
2c00a5a8 XL |
276 | impl<'data, T: 'data> Drop for SliceDrain<'data, T> { |
277 | fn drop(&mut self) { | |
5869c6ff XL |
278 | // extract the iterator so we can use `Drop for [T]` |
279 | let iter = mem::replace(&mut self.iter, [].iter_mut()); | |
280 | unsafe { ptr::drop_in_place(iter.into_slice()) }; | |
2c00a5a8 XL |
281 | } |
282 | } |