specialize StepBy<Range(Inclusive)>
the originally generated code was highly suboptimal this brings it close to the same code or even exactly the same as a manual while-loop by eliminating a branch and the double stepping of n-1 + 1 steps The intermediate trait lets us circumvent the specialization type inference bugs
This commit is contained in:
parent
61ba018093
commit
000aff604e
@ -319,9 +319,10 @@
|
||||
use cmp;
|
||||
use fmt;
|
||||
use iter_private::TrustedRandomAccess;
|
||||
use ops::Try;
|
||||
use ops::{self, Try};
|
||||
use usize;
|
||||
use intrinsics;
|
||||
use mem;
|
||||
|
||||
#[stable(feature = "rust1", since = "1.0.0")]
|
||||
pub use self::iterator::Iterator;
|
||||
@ -672,12 +673,7 @@ impl<I> Iterator for StepBy<I> where I: Iterator {
|
||||
|
||||
#[inline]
|
||||
fn next(&mut self) -> Option<Self::Item> {
|
||||
if self.first_take {
|
||||
self.first_take = false;
|
||||
self.iter.next()
|
||||
} else {
|
||||
self.iter.nth(self.step)
|
||||
}
|
||||
<Self as StepBySpecIterator>::spec_next(self)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
@ -737,6 +733,76 @@ impl<I> Iterator for StepBy<I> where I: Iterator {
|
||||
}
|
||||
}
|
||||
|
||||
// hidden trait for specializing iterator methods
|
||||
// could be generalized but is currently only used for StepBy
|
||||
trait StepBySpecIterator {
|
||||
type Item;
|
||||
fn spec_next(&mut self) -> Option<Self::Item>;
|
||||
}
|
||||
|
||||
impl<I> StepBySpecIterator for StepBy<I>
|
||||
where
|
||||
I: Iterator,
|
||||
{
|
||||
type Item = I::Item;
|
||||
|
||||
#[inline]
|
||||
default fn spec_next(&mut self) -> Option<I::Item> {
|
||||
if self.first_take {
|
||||
self.first_take = false;
|
||||
self.iter.next()
|
||||
} else {
|
||||
self.iter.nth(self.step)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> StepBySpecIterator for StepBy<ops::Range<T>>
|
||||
where
|
||||
T: Step,
|
||||
{
|
||||
#[inline]
|
||||
fn spec_next(&mut self) -> Option<Self::Item> {
|
||||
self.first_take = false;
|
||||
if !(self.iter.start < self.iter.end) {
|
||||
return None;
|
||||
}
|
||||
// add 1 to self.step to get original step size back
|
||||
// it was decremented for the general case on construction
|
||||
if let Some(n) = self.iter.start.add_usize(self.step+1) {
|
||||
let next = mem::replace(&mut self.iter.start, n);
|
||||
Some(next)
|
||||
} else {
|
||||
let last = self.iter.start.clone();
|
||||
self.iter.start = self.iter.end.clone();
|
||||
Some(last)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> StepBySpecIterator for StepBy<ops::RangeInclusive<T>>
|
||||
where
|
||||
T: Step,
|
||||
{
|
||||
#[inline]
|
||||
fn spec_next(&mut self) -> Option<Self::Item> {
|
||||
self.first_take = false;
|
||||
if !(self.iter.start <= self.iter.end) {
|
||||
return None;
|
||||
}
|
||||
// add 1 to self.step to get original step size back
|
||||
// it was decremented for the general case on construction
|
||||
if let Some(n) = self.iter.start.add_usize(self.step+1) {
|
||||
let next = mem::replace(&mut self.iter.start, n);
|
||||
Some(next)
|
||||
} else {
|
||||
let last = self.iter.start.replace_one();
|
||||
self.iter.end.replace_zero();
|
||||
Some(last)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// StepBy can only make the iterator shorter, so the len will still fit.
|
||||
#[stable(feature = "iterator_step_by", since = "1.28.0")]
|
||||
impl<I> ExactSizeIterator for StepBy<I> where I: ExactSizeIterator {}
|
||||
|
@ -1618,6 +1618,14 @@ fn test_range_step() {
|
||||
assert_eq!((isize::MIN..isize::MAX).step_by(1).size_hint(), (usize::MAX, Some(usize::MAX)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_range_inclusive_step() {
|
||||
assert_eq!((0..=50).step_by(10).collect::<Vec<_>>(), [0, 10, 20, 30, 40, 50]);
|
||||
assert_eq!((0..=5).step_by(1).collect::<Vec<_>>(), [0, 1, 2, 3, 4, 5]);
|
||||
assert_eq!((200..=255u8).step_by(10).collect::<Vec<_>>(), [200, 210, 220, 230, 240, 250]);
|
||||
assert_eq!((250..=255u8).step_by(1).collect::<Vec<_>>(), [250, 251, 252, 253, 254, 255]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_range_last_max() {
|
||||
assert_eq!((0..20).last(), Some(19));
|
||||
|
Loading…
Reference in New Issue
Block a user