1 | //! Parallel iterator types for [ranges][std::range], |
2 | //! the type for values created by `a..b` expressions |
3 | //! |
4 | //! You will rarely need to interact with this module directly unless you have |
5 | //! need to name one of the iterator types. |
6 | //! |
7 | //! ``` |
8 | //! use rayon::prelude::*; |
9 | //! |
10 | //! let r = (0..100u64).into_par_iter() |
11 | //! .sum(); |
12 | //! |
13 | //! // compare result with sequential calculation |
14 | //! assert_eq!((0..100).sum::<u64>(), r); |
15 | //! ``` |
16 | //! |
17 | //! [std::range]: https://doc.rust-lang.org/core/ops/struct.Range.html |
18 | |
19 | use crate::iter::plumbing::*; |
20 | use crate::iter::*; |
21 | use std::char; |
22 | use std::convert::TryFrom; |
23 | use std::ops::Range; |
24 | use std::usize; |
25 | |
26 | /// Parallel iterator over a range, implemented for all integer types and `char`. |
27 | /// |
28 | /// **Note:** The `zip` operation requires `IndexedParallelIterator` |
29 | /// which is not implemented for `u64`, `i64`, `u128`, or `i128`. |
30 | /// |
31 | /// ``` |
32 | /// use rayon::prelude::*; |
33 | /// |
34 | /// let p = (0..25usize).into_par_iter() |
35 | /// .zip(0..25usize) |
36 | /// .filter(|&(x, y)| x % 5 == 0 || y % 5 == 0) |
37 | /// .map(|(x, y)| x * y) |
38 | /// .sum::<usize>(); |
39 | /// |
40 | /// let s = (0..25usize).zip(0..25) |
41 | /// .filter(|&(x, y)| x % 5 == 0 || y % 5 == 0) |
42 | /// .map(|(x, y)| x * y) |
43 | /// .sum(); |
44 | /// |
45 | /// assert_eq!(p, s); |
46 | /// ``` |
47 | #[derive(Debug, Clone)] |
48 | pub struct Iter<T> { |
49 | range: Range<T>, |
50 | } |
51 | |
52 | /// Implemented for ranges of all primitive integer types and `char`. |
53 | impl<T> IntoParallelIterator for Range<T> |
54 | where |
55 | Iter<T>: ParallelIterator, |
56 | { |
57 | type Item = <Iter<T> as ParallelIterator>::Item; |
58 | type Iter = Iter<T>; |
59 | |
60 | fn into_par_iter(self) -> Self::Iter { |
61 | Iter { range: self } |
62 | } |
63 | } |
64 | |
65 | struct IterProducer<T> { |
66 | range: Range<T>, |
67 | } |
68 | |
69 | impl<T> IntoIterator for IterProducer<T> |
70 | where |
71 | Range<T>: Iterator, |
72 | { |
73 | type Item = <Range<T> as Iterator>::Item; |
74 | type IntoIter = Range<T>; |
75 | |
76 | fn into_iter(self) -> Self::IntoIter { |
77 | self.range |
78 | } |
79 | } |
80 | |
81 | /// These traits help drive integer type inference. Without them, an unknown `{integer}` type only |
82 | /// has constraints on `Iter<{integer}>`, which will probably give up and use `i32`. By adding |
83 | /// these traits on the item type, the compiler can see a more direct constraint to infer like |
84 | /// `{integer}: RangeInteger`, which works better. See `test_issue_833` for an example. |
85 | /// |
86 | /// They have to be `pub` since they're seen in the public `impl ParallelIterator` constraints, but |
87 | /// we put them in a private modules so they're not actually reachable in our public API. |
88 | mod private { |
89 | use super::*; |
90 | |
91 | /// Implementation details of `ParallelIterator for Iter<Self>` |
92 | pub trait RangeInteger: Sized + Send { |
93 | private_decl! {} |
94 | |
95 | fn drive_unindexed<C>(iter: Iter<Self>, consumer: C) -> C::Result |
96 | where |
97 | C: UnindexedConsumer<Self>; |
98 | |
99 | fn opt_len(iter: &Iter<Self>) -> Option<usize>; |
100 | } |
101 | |
102 | /// Implementation details of `IndexedParallelIterator for Iter<Self>` |
103 | pub trait IndexedRangeInteger: RangeInteger { |
104 | private_decl! {} |
105 | |
106 | fn drive<C>(iter: Iter<Self>, consumer: C) -> C::Result |
107 | where |
108 | C: Consumer<Self>; |
109 | |
110 | fn len(iter: &Iter<Self>) -> usize; |
111 | |
112 | fn with_producer<CB>(iter: Iter<Self>, callback: CB) -> CB::Output |
113 | where |
114 | CB: ProducerCallback<Self>; |
115 | } |
116 | } |
117 | use private::{IndexedRangeInteger, RangeInteger}; |
118 | |
119 | impl<T: RangeInteger> ParallelIterator for Iter<T> { |
120 | type Item = T; |
121 | |
122 | fn drive_unindexed<C>(self, consumer: C) -> C::Result |
123 | where |
124 | C: UnindexedConsumer<T>, |
125 | { |
126 | T::drive_unindexed(self, consumer) |
127 | } |
128 | |
129 | #[inline ] |
130 | fn opt_len(&self) -> Option<usize> { |
131 | T::opt_len(self) |
132 | } |
133 | } |
134 | |
135 | impl<T: IndexedRangeInteger> IndexedParallelIterator for Iter<T> { |
136 | fn drive<C>(self, consumer: C) -> C::Result |
137 | where |
138 | C: Consumer<T>, |
139 | { |
140 | T::drive(self, consumer) |
141 | } |
142 | |
143 | #[inline ] |
144 | fn len(&self) -> usize { |
145 | T::len(self) |
146 | } |
147 | |
148 | fn with_producer<CB>(self, callback: CB) -> CB::Output |
149 | where |
150 | CB: ProducerCallback<T>, |
151 | { |
152 | T::with_producer(self, callback) |
153 | } |
154 | } |
155 | |
156 | macro_rules! indexed_range_impl { |
157 | ( $t:ty ) => { |
158 | impl RangeInteger for $t { |
159 | private_impl! {} |
160 | |
161 | fn drive_unindexed<C>(iter: Iter<$t>, consumer: C) -> C::Result |
162 | where |
163 | C: UnindexedConsumer<$t>, |
164 | { |
165 | bridge(iter, consumer) |
166 | } |
167 | |
168 | fn opt_len(iter: &Iter<$t>) -> Option<usize> { |
169 | Some(iter.range.len()) |
170 | } |
171 | } |
172 | |
173 | impl IndexedRangeInteger for $t { |
174 | private_impl! {} |
175 | |
176 | fn drive<C>(iter: Iter<$t>, consumer: C) -> C::Result |
177 | where |
178 | C: Consumer<$t>, |
179 | { |
180 | bridge(iter, consumer) |
181 | } |
182 | |
183 | fn len(iter: &Iter<$t>) -> usize { |
184 | iter.range.len() |
185 | } |
186 | |
187 | fn with_producer<CB>(iter: Iter<$t>, callback: CB) -> CB::Output |
188 | where |
189 | CB: ProducerCallback<$t>, |
190 | { |
191 | callback.callback(IterProducer { range: iter.range }) |
192 | } |
193 | } |
194 | |
195 | impl Producer for IterProducer<$t> { |
196 | type Item = <Range<$t> as Iterator>::Item; |
197 | type IntoIter = Range<$t>; |
198 | fn into_iter(self) -> Self::IntoIter { |
199 | self.range |
200 | } |
201 | |
202 | fn split_at(self, index: usize) -> (Self, Self) { |
203 | assert!(index <= self.range.len()); |
204 | // For signed $t, the length and requested index could be greater than $t::MAX, and |
205 | // then `index as $t` could wrap to negative, so wrapping_add is necessary. |
206 | let mid = self.range.start.wrapping_add(index as $t); |
207 | let left = self.range.start..mid; |
208 | let right = mid..self.range.end; |
209 | (IterProducer { range: left }, IterProducer { range: right }) |
210 | } |
211 | } |
212 | }; |
213 | } |
214 | |
215 | trait UnindexedRangeLen<L> { |
216 | fn len(&self) -> L; |
217 | } |
218 | |
219 | macro_rules! unindexed_range_impl { |
220 | ( $t:ty, $len_t:ty ) => { |
221 | impl UnindexedRangeLen<$len_t> for Range<$t> { |
222 | fn len(&self) -> $len_t { |
223 | let &Range { start, end } = self; |
224 | if end > start { |
225 | end.wrapping_sub(start) as $len_t |
226 | } else { |
227 | 0 |
228 | } |
229 | } |
230 | } |
231 | |
232 | impl RangeInteger for $t { |
233 | private_impl! {} |
234 | |
235 | fn drive_unindexed<C>(iter: Iter<$t>, consumer: C) -> C::Result |
236 | where |
237 | C: UnindexedConsumer<$t>, |
238 | { |
239 | #[inline] |
240 | fn offset(start: $t) -> impl Fn(usize) -> $t { |
241 | move |i| start.wrapping_add(i as $t) |
242 | } |
243 | |
244 | if let Some(len) = iter.opt_len() { |
245 | // Drive this in indexed mode for better `collect`. |
246 | (0..len) |
247 | .into_par_iter() |
248 | .map(offset(iter.range.start)) |
249 | .drive(consumer) |
250 | } else { |
251 | bridge_unindexed(IterProducer { range: iter.range }, consumer) |
252 | } |
253 | } |
254 | |
255 | fn opt_len(iter: &Iter<$t>) -> Option<usize> { |
256 | usize::try_from(iter.range.len()).ok() |
257 | } |
258 | } |
259 | |
260 | impl UnindexedProducer for IterProducer<$t> { |
261 | type Item = $t; |
262 | |
263 | fn split(mut self) -> (Self, Option<Self>) { |
264 | let index = self.range.len() / 2; |
265 | if index > 0 { |
266 | let mid = self.range.start.wrapping_add(index as $t); |
267 | let right = mid..self.range.end; |
268 | self.range.end = mid; |
269 | (self, Some(IterProducer { range: right })) |
270 | } else { |
271 | (self, None) |
272 | } |
273 | } |
274 | |
275 | fn fold_with<F>(self, folder: F) -> F |
276 | where |
277 | F: Folder<Self::Item>, |
278 | { |
279 | folder.consume_iter(self) |
280 | } |
281 | } |
282 | }; |
283 | } |
284 | |
285 | // all Range<T> with ExactSizeIterator |
286 | indexed_range_impl! {u8} |
287 | indexed_range_impl! {u16} |
288 | indexed_range_impl! {u32} |
289 | indexed_range_impl! {usize} |
290 | indexed_range_impl! {i8} |
291 | indexed_range_impl! {i16} |
292 | indexed_range_impl! {i32} |
293 | indexed_range_impl! {isize} |
294 | |
295 | // other Range<T> with just Iterator |
296 | unindexed_range_impl! {u64, u64} |
297 | unindexed_range_impl! {i64, u64} |
298 | unindexed_range_impl! {u128, u128} |
299 | unindexed_range_impl! {i128, u128} |
300 | |
301 | // char is special because of the surrogate range hole |
302 | macro_rules! convert_char { |
303 | ( $self:ident . $method:ident ( $( $arg:expr ),* ) ) => {{ |
304 | let start = $self.range.start as u32; |
305 | let end = $self.range.end as u32; |
306 | if start < 0xD800 && 0xE000 < end { |
307 | // chain the before and after surrogate range fragments |
308 | (start..0xD800) |
309 | .into_par_iter() |
310 | .chain(0xE000..end) |
311 | .map(|codepoint| unsafe { char::from_u32_unchecked(codepoint) }) |
312 | .$method($( $arg ),*) |
313 | } else { |
314 | // no surrogate range to worry about |
315 | (start..end) |
316 | .into_par_iter() |
317 | .map(|codepoint| unsafe { char::from_u32_unchecked(codepoint) }) |
318 | .$method($( $arg ),*) |
319 | } |
320 | }}; |
321 | } |
322 | |
323 | impl ParallelIterator for Iter<char> { |
324 | type Item = char; |
325 | |
326 | fn drive_unindexed<C>(self, consumer: C) -> C::Result |
327 | where |
328 | C: UnindexedConsumer<Self::Item>, |
329 | { |
330 | convert_char!(self.drive(consumer)) |
331 | } |
332 | |
333 | fn opt_len(&self) -> Option<usize> { |
334 | Some(self.len()) |
335 | } |
336 | } |
337 | |
338 | impl IndexedParallelIterator for Iter<char> { |
339 | // Split at the surrogate range first if we're allowed to |
340 | fn drive<C>(self, consumer: C) -> C::Result |
341 | where |
342 | C: Consumer<Self::Item>, |
343 | { |
344 | convert_char!(self.drive(consumer)) |
345 | } |
346 | |
347 | fn len(&self) -> usize { |
348 | // Taken from <char as Step>::steps_between |
349 | let start = self.range.start as u32; |
350 | let end = self.range.end as u32; |
351 | if start < end { |
352 | let mut count = end - start; |
353 | if start < 0xD800 && 0xE000 <= end { |
354 | count -= 0x800 |
355 | } |
356 | count as usize |
357 | } else { |
358 | 0 |
359 | } |
360 | } |
361 | |
362 | fn with_producer<CB>(self, callback: CB) -> CB::Output |
363 | where |
364 | CB: ProducerCallback<Self::Item>, |
365 | { |
366 | convert_char!(self.with_producer(callback)) |
367 | } |
368 | } |
369 | |
370 | #[test] |
371 | fn check_range_split_at_overflow() { |
372 | // Note, this split index overflows i8! |
373 | let producer = IterProducer { range: -100i8..100 }; |
374 | let (left, right) = producer.split_at(150); |
375 | let r1: i32 = left.range.map(i32::from).sum(); |
376 | let r2: i32 = right.range.map(i32::from).sum(); |
377 | assert_eq!(r1 + r2, -100); |
378 | } |
379 | |
380 | #[test] |
381 | fn test_i128_len_doesnt_overflow() { |
382 | use std::{i128, u128}; |
383 | |
384 | // Using parse because some versions of rust don't allow long literals |
385 | let octillion: i128 = "1000000000000000000000000000" .parse().unwrap(); |
386 | let producer = IterProducer { |
387 | range: 0..octillion, |
388 | }; |
389 | |
390 | assert_eq!(octillion as u128, producer.range.len()); |
391 | assert_eq!(octillion as u128, (0..octillion).len()); |
392 | assert_eq!(2 * octillion as u128, (-octillion..octillion).len()); |
393 | |
394 | assert_eq!(u128::MAX, (i128::MIN..i128::MAX).len()); |
395 | } |
396 | |
397 | #[test] |
398 | fn test_u64_opt_len() { |
399 | use std::{u64, usize}; |
400 | assert_eq!(Some(100), (0..100u64).into_par_iter().opt_len()); |
401 | assert_eq!( |
402 | Some(usize::MAX), |
403 | (0..usize::MAX as u64).into_par_iter().opt_len() |
404 | ); |
405 | if (usize::MAX as u64) < u64::MAX { |
406 | assert_eq!( |
407 | None, |
408 | (0..(usize::MAX as u64).wrapping_add(1)) |
409 | .into_par_iter() |
410 | .opt_len() |
411 | ); |
412 | assert_eq!(None, (0..u64::MAX).into_par_iter().opt_len()); |
413 | } |
414 | } |
415 | |
416 | #[test] |
417 | fn test_u128_opt_len() { |
418 | use std::{u128, usize}; |
419 | assert_eq!(Some(100), (0..100u128).into_par_iter().opt_len()); |
420 | assert_eq!( |
421 | Some(usize::MAX), |
422 | (0..usize::MAX as u128).into_par_iter().opt_len() |
423 | ); |
424 | assert_eq!(None, (0..1 + usize::MAX as u128).into_par_iter().opt_len()); |
425 | assert_eq!(None, (0..u128::MAX).into_par_iter().opt_len()); |
426 | } |
427 | |
428 | // `usize as i64` can overflow, so make sure to wrap it appropriately |
429 | // when using the `opt_len` "indexed" mode. |
430 | #[test] |
431 | #[cfg (target_pointer_width = "64" )] |
432 | fn test_usize_i64_overflow() { |
433 | use crate::ThreadPoolBuilder; |
434 | use std::i64; |
435 | |
436 | let iter = (-2..i64::MAX).into_par_iter(); |
437 | assert_eq!(iter.opt_len(), Some(i64::MAX as usize + 2)); |
438 | |
439 | // always run with multiple threads to split into, or this will take forever... |
440 | let pool = ThreadPoolBuilder::new().num_threads(8).build().unwrap(); |
441 | pool.install(|| assert_eq!(iter.find_last(|_| true), Some(i64::MAX - 1))); |
442 | } |
443 | |
444 | #[test] |
445 | fn test_issue_833() { |
446 | fn is_even(n: i64) -> bool { |
447 | n % 2 == 0 |
448 | } |
449 | |
450 | // The integer type should be inferred from `is_even` |
451 | let v: Vec<_> = (1..100).into_par_iter().filter(|&x| is_even(x)).collect(); |
452 | assert!(v.into_iter().eq((2..100).step_by(2))); |
453 | |
454 | // Try examples with indexed iterators too |
455 | let pos = (0..100).into_par_iter().position_any(|x| x == 50i16); |
456 | assert_eq!(pos, Some(50usize)); |
457 | |
458 | assert!((0..100) |
459 | .into_par_iter() |
460 | .zip(0..100) |
461 | .all(|(a, b)| i16::eq(&a, &b))); |
462 | } |
463 | |