about summary refs log tree commit diff
diff options
context:
space:
mode:
authorEmerentius <emerentius@arcor.de>2018-06-16 21:58:28 +0200
committerEmerentius <emerentius@arcor.de>2018-06-19 19:33:54 +0200
commit000aff604e3b16ffc3771bd5d93a6e7b425852d2 (patch)
tree3bd5ca2639264c835fe5e39bea26d95a687e475a
parent61ba0180933485cf8a2bc6b7230a4c70b82bb063 (diff)
downloadrust-000aff604e3b16ffc3771bd5d93a6e7b425852d2.tar.gz
rust-000aff604e3b16ffc3771bd5d93a6e7b425852d2.zip
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
-rw-r--r--src/libcore/iter/mod.rs80
-rw-r--r--src/libcore/tests/iter.rs8
2 files changed, 81 insertions, 7 deletions
diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs
index 840d45ff1cc..c4132270d59 100644
--- a/src/libcore/iter/mod.rs
+++ b/src/libcore/iter/mod.rs
@@ -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 {}
diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs
index 9b8d7031f8e..72b115f8b5f 100644
--- a/src/libcore/tests/iter.rs
+++ b/src/libcore/tests/iter.rs
@@ -1619,6 +1619,14 @@ fn test_range_step() {
 }
 
 #[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));
     assert_eq!((-20..0).last(), Some(-1));