Skip to content

Commit 383af35

Browse files
committed
Implement opt_len for ranges of large integers
If the length of a range of `i64`, `u64`, `i128`, or `u128` would actually fit in `usize`, we can return that from `opt_len` to enable some optimizations, namely indexed `collect` into vectors. To drive such ranges, we'll actually iterate over `(0..len)` instead, and then map it to the larger type.
1 parent da16c2c commit 383af35

File tree

2 files changed

+69
-6
lines changed

2 files changed

+69
-6
lines changed

src/iter/test.rs

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,21 @@ pub fn execute_unindexed_range() {
5050
assert_eq!(b, c);
5151
}
5252

53+
#[cfg(has_i128)]
54+
#[test]
55+
pub fn execute_pseudo_indexed_range() {
56+
use std::i128::MAX;
57+
let range = MAX - 1024..MAX;
58+
59+
// Given `Some` length, collecting `Vec` will try to act indexed.
60+
let a = range.clone().into_par_iter();
61+
assert_eq!(a.opt_len(), Some(1024));
62+
63+
let b: Vec<i128> = a.map(|i| i + 1).collect();
64+
let c: Vec<i128> = range.map(|i| i + 1).collect();
65+
assert_eq!(b, c);
66+
}
67+
5368
#[test]
5469
pub fn check_map_indexed() {
5570
let a = [1, 2, 3];

src/range.rs

Lines changed: 54 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
use iter::*;
2020
use iter::plumbing::*;
2121
use std::ops::Range;
22+
use std::usize;
2223

2324
/// Parallel iterator over a range, implemented for all integer types.
2425
///
@@ -127,11 +128,15 @@ macro_rules! indexed_range_impl {
127128
}
128129
}
129130

131+
trait UnindexedRangeLen<L> {
132+
fn len(&self) -> L;
133+
}
134+
130135
macro_rules! unindexed_range_impl {
131136
( $t:ty, $len_t:ty ) => {
132-
impl IterProducer<$t> {
137+
impl UnindexedRangeLen<$len_t> for Range<$t> {
133138
fn len(&self) -> $len_t {
134-
let Range { start, end } = self.range;
139+
let &Range { start, end } = self;
135140
if end > start {
136141
end.wrapping_sub(start) as $len_t
137142
} else {
@@ -146,15 +151,30 @@ macro_rules! unindexed_range_impl {
146151
fn drive_unindexed<C>(self, consumer: C) -> C::Result
147152
where C: UnindexedConsumer<Self::Item>
148153
{
149-
bridge_unindexed(IterProducer { range: self.range }, consumer)
154+
if let Some(len) = self.opt_len() {
155+
(0..len).into_par_iter()
156+
.map(|i| self.range.start + i as $t)
157+
.drive_unindexed(consumer)
158+
} else {
159+
bridge_unindexed(IterProducer { range: self.range }, consumer)
160+
}
161+
}
162+
163+
fn opt_len(&self) -> Option<usize> {
164+
let len = self.range.len();
165+
if len <= usize::MAX as $len_t {
166+
Some(len as usize)
167+
} else {
168+
None
169+
}
150170
}
151171
}
152172

153173
impl UnindexedProducer for IterProducer<$t> {
154174
type Item = $t;
155175

156176
fn split(mut self) -> (Self, Option<Self>) {
157-
let index = self.len() / 2;
177+
let index = self.range.len() / 2;
158178
if index > 0 {
159179
let mid = self.range.start.wrapping_add(index as $t);
160180
let right = mid .. self.range.end;
@@ -203,9 +223,37 @@ pub fn check_range_split_at_overflow() {
203223
#[cfg(has_i128)]
204224
#[test]
205225
pub fn test_i128_len_doesnt_overflow() {
226+
use std::{i128, u128};
227+
206228
// Using parse because some versions of rust don't allow long literals
207-
let octillion = "1000000000000000000000000000".parse::<i128>().unwrap();
229+
let octillion: i128 = "1000000000000000000000000000".parse().unwrap();
208230
let producer = IterProducer { range: 0..octillion };
209231

210-
assert_eq!(octillion as u128, producer.len());
232+
assert_eq!(octillion as u128, producer.range.len());
233+
assert_eq!(octillion as u128, (0..octillion).len());
234+
assert_eq!(2 * octillion as u128, (-octillion..octillion).len());
235+
236+
assert_eq!(u128::MAX, (i128::MIN..i128::MAX).len());
237+
238+
}
239+
240+
#[test]
241+
pub fn test_u64_opt_len() {
242+
use std::{u64, usize};
243+
assert_eq!(Some(100), (0..100u64).into_par_iter().opt_len());
244+
assert_eq!(Some(usize::MAX), (0..usize::MAX as u64).into_par_iter().opt_len());
245+
if (usize::MAX as u64) < u64::MAX {
246+
assert_eq!(None, (0..1 + usize::MAX as u64).into_par_iter().opt_len());
247+
assert_eq!(None, (0..u64::MAX).into_par_iter().opt_len());
248+
}
249+
}
250+
251+
#[cfg(has_i128)]
252+
#[test]
253+
pub fn test_u128_opt_len() {
254+
use std::{u128, usize};
255+
assert_eq!(Some(100), (0..100u128).into_par_iter().opt_len());
256+
assert_eq!(Some(usize::MAX), (0..usize::MAX as u128).into_par_iter().opt_len());
257+
assert_eq!(None, (0..1 + usize::MAX as u128).into_par_iter().opt_len());
258+
assert_eq!(None, (0..u128::MAX).into_par_iter().opt_len());
211259
}

0 commit comments

Comments
 (0)