1 | //! Parallel iterator types for [inclusive 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.RangeInclusive.html |
18 | |
19 | use crate::iter::plumbing::*; |
20 | use crate::iter::*; |
21 | use std::char; |
22 | use std::ops::RangeInclusive; |
23 | |
24 | /// Parallel iterator over an inclusive range, implemented for all integer types and `char`. |
25 | /// |
26 | /// **Note:** The `zip` operation requires `IndexedParallelIterator` |
27 | /// which is only implemented for `u8`, `i8`, `u16`, `i16`, and `char`. |
28 | /// |
29 | /// ``` |
30 | /// use rayon::prelude::*; |
31 | /// |
32 | /// let p = (0..=25u16).into_par_iter() |
33 | /// .zip(0..=25u16) |
34 | /// .filter(|&(x, y)| x % 5 == 0 || y % 5 == 0) |
35 | /// .map(|(x, y)| x * y) |
36 | /// .sum::<u16>(); |
37 | /// |
38 | /// let s = (0..=25u16).zip(0..=25u16) |
39 | /// .filter(|&(x, y)| x % 5 == 0 || y % 5 == 0) |
40 | /// .map(|(x, y)| x * y) |
41 | /// .sum(); |
42 | /// |
43 | /// assert_eq!(p, s); |
44 | /// ``` |
45 | #[derive(Debug, Clone)] |
46 | pub struct Iter<T> { |
47 | range: RangeInclusive<T>, |
48 | } |
49 | |
50 | impl<T> Iter<T> |
51 | where |
52 | RangeInclusive<T>: Eq, |
53 | T: Ord + Copy, |
54 | { |
55 | /// Returns `Some((start, end))` for `start..=end`, or `None` if it is exhausted. |
56 | /// |
57 | /// Note that `RangeInclusive` does not specify the bounds of an exhausted iterator, |
58 | /// so this is a way for us to figure out what we've got. Thankfully, all of the |
59 | /// integer types we care about can be trivially cloned. |
60 | fn bounds(&self) -> Option<(T, T)> { |
61 | let start = *self.range.start(); |
62 | let end = *self.range.end(); |
63 | if start <= end && self.range == (start..=end) { |
64 | // If the range is still nonempty, this is obviously true |
65 | // If the range is exhausted, either start > end or |
66 | // the range does not equal start..=end. |
67 | Some((start, end)) |
68 | } else { |
69 | None |
70 | } |
71 | } |
72 | } |
73 | |
74 | /// Implemented for ranges of all primitive integer types and `char`. |
75 | impl<T> IntoParallelIterator for RangeInclusive<T> |
76 | where |
77 | Iter<T>: ParallelIterator, |
78 | { |
79 | type Item = <Iter<T> as ParallelIterator>::Item; |
80 | type Iter = Iter<T>; |
81 | |
82 | fn into_par_iter(self) -> Self::Iter { |
83 | Iter { range: self } |
84 | } |
85 | } |
86 | |
87 | /// These traits help drive integer type inference. Without them, an unknown `{integer}` type only |
88 | /// has constraints on `Iter<{integer}>`, which will probably give up and use `i32`. By adding |
89 | /// these traits on the item type, the compiler can see a more direct constraint to infer like |
90 | /// `{integer}: RangeInteger`, which works better. See `test_issue_833` for an example. |
91 | /// |
92 | /// They have to be `pub` since they're seen in the public `impl ParallelIterator` constraints, but |
93 | /// we put them in a private modules so they're not actually reachable in our public API. |
94 | mod private { |
95 | use super::*; |
96 | |
97 | /// Implementation details of `ParallelIterator for Iter<Self>` |
98 | pub trait RangeInteger: Sized + Send { |
99 | private_decl! {} |
100 | |
101 | fn drive_unindexed<C>(iter: Iter<Self>, consumer: C) -> C::Result |
102 | where |
103 | C: UnindexedConsumer<Self>; |
104 | |
105 | fn opt_len(iter: &Iter<Self>) -> Option<usize>; |
106 | } |
107 | |
108 | /// Implementation details of `IndexedParallelIterator for Iter<Self>` |
109 | pub trait IndexedRangeInteger: RangeInteger { |
110 | private_decl! {} |
111 | |
112 | fn drive<C>(iter: Iter<Self>, consumer: C) -> C::Result |
113 | where |
114 | C: Consumer<Self>; |
115 | |
116 | fn len(iter: &Iter<Self>) -> usize; |
117 | |
118 | fn with_producer<CB>(iter: Iter<Self>, callback: CB) -> CB::Output |
119 | where |
120 | CB: ProducerCallback<Self>; |
121 | } |
122 | } |
123 | use private::{IndexedRangeInteger, RangeInteger}; |
124 | |
125 | impl<T: RangeInteger> ParallelIterator for Iter<T> { |
126 | type Item = T; |
127 | |
128 | fn drive_unindexed<C>(self, consumer: C) -> C::Result |
129 | where |
130 | C: UnindexedConsumer<T>, |
131 | { |
132 | T::drive_unindexed(self, consumer) |
133 | } |
134 | |
135 | #[inline ] |
136 | fn opt_len(&self) -> Option<usize> { |
137 | T::opt_len(self) |
138 | } |
139 | } |
140 | |
141 | impl<T: IndexedRangeInteger> IndexedParallelIterator for Iter<T> { |
142 | fn drive<C>(self, consumer: C) -> C::Result |
143 | where |
144 | C: Consumer<T>, |
145 | { |
146 | T::drive(self, consumer) |
147 | } |
148 | |
149 | #[inline ] |
150 | fn len(&self) -> usize { |
151 | T::len(self) |
152 | } |
153 | |
154 | fn with_producer<CB>(self, callback: CB) -> CB::Output |
155 | where |
156 | CB: ProducerCallback<T>, |
157 | { |
158 | T::with_producer(self, callback) |
159 | } |
160 | } |
161 | |
162 | macro_rules! convert { |
163 | ( $iter:ident . $method:ident ( $( $arg:expr ),* ) ) => { |
164 | if let Some((start, end)) = $iter.bounds() { |
165 | if let Some(end) = end.checked_add(1) { |
166 | (start..end).into_par_iter().$method($( $arg ),*) |
167 | } else { |
168 | (start..end).into_par_iter().chain(once(end)).$method($( $arg ),*) |
169 | } |
170 | } else { |
171 | empty::<Self>().$method($( $arg ),*) |
172 | } |
173 | }; |
174 | } |
175 | |
176 | macro_rules! parallel_range_impl { |
177 | ( $t:ty ) => { |
178 | impl RangeInteger for $t { |
179 | private_impl! {} |
180 | |
181 | fn drive_unindexed<C>(iter: Iter<$t>, consumer: C) -> C::Result |
182 | where |
183 | C: UnindexedConsumer<$t>, |
184 | { |
185 | convert!(iter.drive_unindexed(consumer)) |
186 | } |
187 | |
188 | fn opt_len(iter: &Iter<$t>) -> Option<usize> { |
189 | convert!(iter.opt_len()) |
190 | } |
191 | } |
192 | }; |
193 | } |
194 | |
195 | macro_rules! indexed_range_impl { |
196 | ( $t:ty ) => { |
197 | parallel_range_impl! { $t } |
198 | |
199 | impl IndexedRangeInteger for $t { |
200 | private_impl! {} |
201 | |
202 | fn drive<C>(iter: Iter<$t>, consumer: C) -> C::Result |
203 | where |
204 | C: Consumer<$t>, |
205 | { |
206 | convert!(iter.drive(consumer)) |
207 | } |
208 | |
209 | fn len(iter: &Iter<$t>) -> usize { |
210 | iter.range.len() |
211 | } |
212 | |
213 | fn with_producer<CB>(iter: Iter<$t>, callback: CB) -> CB::Output |
214 | where |
215 | CB: ProducerCallback<$t>, |
216 | { |
217 | convert!(iter.with_producer(callback)) |
218 | } |
219 | } |
220 | }; |
221 | } |
222 | |
223 | // all RangeInclusive<T> with ExactSizeIterator |
224 | indexed_range_impl! {u8} |
225 | indexed_range_impl! {u16} |
226 | indexed_range_impl! {i8} |
227 | indexed_range_impl! {i16} |
228 | |
229 | // other RangeInclusive<T> with just Iterator |
230 | parallel_range_impl! {usize} |
231 | parallel_range_impl! {isize} |
232 | parallel_range_impl! {u32} |
233 | parallel_range_impl! {i32} |
234 | parallel_range_impl! {u64} |
235 | parallel_range_impl! {i64} |
236 | parallel_range_impl! {u128} |
237 | parallel_range_impl! {i128} |
238 | |
239 | // char is special |
240 | macro_rules! convert_char { |
241 | ( $self:ident . $method:ident ( $( $arg:expr ),* ) ) => { |
242 | if let Some((start, end)) = $self.bounds() { |
243 | let start = start as u32; |
244 | let end = end as u32; |
245 | if start < 0xD800 && 0xE000 <= end { |
246 | // chain the before and after surrogate range fragments |
247 | (start..0xD800) |
248 | .into_par_iter() |
249 | .chain(0xE000..end + 1) // cannot use RangeInclusive, so add one to end |
250 | .map(|codepoint| unsafe { char::from_u32_unchecked(codepoint) }) |
251 | .$method($( $arg ),*) |
252 | } else { |
253 | // no surrogate range to worry about |
254 | (start..end + 1) // cannot use RangeInclusive, so add one to end |
255 | .into_par_iter() |
256 | .map(|codepoint| unsafe { char::from_u32_unchecked(codepoint) }) |
257 | .$method($( $arg ),*) |
258 | } |
259 | } else { |
260 | empty::<char>().$method($( $arg ),*) |
261 | } |
262 | }; |
263 | } |
264 | |
265 | impl ParallelIterator for Iter<char> { |
266 | type Item = char; |
267 | |
268 | fn drive_unindexed<C>(self, consumer: C) -> C::Result |
269 | where |
270 | C: UnindexedConsumer<Self::Item>, |
271 | { |
272 | convert_char!(self.drive(consumer)) |
273 | } |
274 | |
275 | fn opt_len(&self) -> Option<usize> { |
276 | Some(self.len()) |
277 | } |
278 | } |
279 | |
280 | // Range<u32> is broken on 16 bit platforms, may as well benefit from it |
281 | impl IndexedParallelIterator for Iter<char> { |
282 | // Split at the surrogate range first if we're allowed to |
283 | fn drive<C>(self, consumer: C) -> C::Result |
284 | where |
285 | C: Consumer<Self::Item>, |
286 | { |
287 | convert_char!(self.drive(consumer)) |
288 | } |
289 | |
290 | fn len(&self) -> usize { |
291 | if let Some((start, end)) = self.bounds() { |
292 | // Taken from <char as Step>::steps_between |
293 | let start = start as u32; |
294 | let end = end as u32; |
295 | let mut count = end - start; |
296 | if start < 0xD800 && 0xE000 <= end { |
297 | count -= 0x800 |
298 | } |
299 | (count + 1) as usize // add one for inclusive |
300 | } else { |
301 | 0 |
302 | } |
303 | } |
304 | |
305 | fn with_producer<CB>(self, callback: CB) -> CB::Output |
306 | where |
307 | CB: ProducerCallback<Self::Item>, |
308 | { |
309 | convert_char!(self.with_producer(callback)) |
310 | } |
311 | } |
312 | |
313 | #[test] |
314 | #[cfg (target_pointer_width = "64" )] |
315 | fn test_u32_opt_len() { |
316 | use std::u32; |
317 | assert_eq!(Some(101), (0..=100u32).into_par_iter().opt_len()); |
318 | assert_eq!( |
319 | Some(u32::MAX as usize), |
320 | (0..=u32::MAX - 1).into_par_iter().opt_len() |
321 | ); |
322 | assert_eq!( |
323 | Some(u32::MAX as usize + 1), |
324 | (0..=u32::MAX).into_par_iter().opt_len() |
325 | ); |
326 | } |
327 | |
328 | #[test] |
329 | fn test_u64_opt_len() { |
330 | use std::{u64, usize}; |
331 | assert_eq!(Some(101), (0..=100u64).into_par_iter().opt_len()); |
332 | assert_eq!( |
333 | Some(usize::MAX), |
334 | (0..=usize::MAX as u64 - 1).into_par_iter().opt_len() |
335 | ); |
336 | assert_eq!(None, (0..=usize::MAX as u64).into_par_iter().opt_len()); |
337 | assert_eq!(None, (0..=u64::MAX).into_par_iter().opt_len()); |
338 | } |
339 | |
340 | #[test] |
341 | fn test_u128_opt_len() { |
342 | use std::{u128, usize}; |
343 | assert_eq!(Some(101), (0..=100u128).into_par_iter().opt_len()); |
344 | assert_eq!( |
345 | Some(usize::MAX), |
346 | (0..=usize::MAX as u128 - 1).into_par_iter().opt_len() |
347 | ); |
348 | assert_eq!(None, (0..=usize::MAX as u128).into_par_iter().opt_len()); |
349 | assert_eq!(None, (0..=u128::MAX).into_par_iter().opt_len()); |
350 | } |
351 | |
352 | // `usize as i64` can overflow, so make sure to wrap it appropriately |
353 | // when using the `opt_len` "indexed" mode. |
354 | #[test] |
355 | #[cfg (target_pointer_width = "64" )] |
356 | fn test_usize_i64_overflow() { |
357 | use crate::ThreadPoolBuilder; |
358 | use std::i64; |
359 | |
360 | let iter = (-2..=i64::MAX).into_par_iter(); |
361 | assert_eq!(iter.opt_len(), Some(i64::MAX as usize + 3)); |
362 | |
363 | // always run with multiple threads to split into, or this will take forever... |
364 | let pool = ThreadPoolBuilder::new().num_threads(8).build().unwrap(); |
365 | pool.install(|| assert_eq!(iter.find_last(|_| true), Some(i64::MAX))); |
366 | } |
367 | |
368 | #[test] |
369 | fn test_issue_833() { |
370 | fn is_even(n: i64) -> bool { |
371 | n % 2 == 0 |
372 | } |
373 | |
374 | // The integer type should be inferred from `is_even` |
375 | let v: Vec<_> = (1..=100).into_par_iter().filter(|&x| is_even(x)).collect(); |
376 | assert!(v.into_iter().eq((2..=100).step_by(2))); |
377 | |
378 | // Try examples with indexed iterators too |
379 | let pos = (0..=100).into_par_iter().position_any(|x| x == 50i16); |
380 | assert_eq!(pos, Some(50usize)); |
381 | |
382 | assert!((0..=100) |
383 | .into_par_iter() |
384 | .zip(0..=100) |
385 | .all(|(a, b)| i16::eq(&a, &b))); |
386 | } |
387 | |