Override size_hint and propagate ExactSizeIterator for iter::StepBy

Generally useful, but also a prerequisite for moving a bunch of unit tests off Range::step_by.

3 files changed, 94 insertions, 0 deletions

diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs
index 420ff0f7119..a3b755a38f3 100644
--- a/src/libcore/iter/mod.rs
+++ b/src/libcore/iter/mod.rs
@@ -553,8 +553,27 @@ impl<I> Iterator for StepBy<I> where I: Iterator {
             self.iter.nth(self.step)
         }
     }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let inner_hint = self.iter.size_hint();
+
+        if self.first_take {
+            let f = |n| if n == 0 { 0 } else { 1 + (n-1)/(self.step+1) };
+            (f(inner_hint.0), inner_hint.1.map(f))
+        } else {
+            let f = |n| n / (self.step+1);
+            (f(inner_hint.0), inner_hint.1.map(f))
+        }
+    }
 }
 
+// StepBy can only make the iterator shorter, so the len will still fit.
+#[unstable(feature = "iterator_step_by",
+           reason = "unstable replacement of Range::step_by",
+           issue = "27741")]
+impl<I> ExactSizeIterator for StepBy<I> where I: ExactSizeIterator {}
+
 /// An iterator that strings two iterators together.
 ///
 /// This `struct` is created by the [`chain`] method on [`Iterator`]. See its
diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs
index ad91ba9be58..4030eaf2b23 100644
--- a/src/libcore/tests/iter.rs
+++ b/src/libcore/tests/iter.rs
@@ -172,6 +172,79 @@ fn test_iterator_step_by_zero() {
 }
 
 #[test]
+fn test_iterator_step_by_size_hint() {
+    struct StubSizeHint(usize, Option<usize>);
+    impl Iterator for StubSizeHint {
+        type Item = ();
+        fn next(&mut self) -> Option<()> {
+            self.0 -= 1;
+            if let Some(ref mut upper) = self.1 {
+                *upper -= 1;
+            }
+            Some(())
+        }
+        fn size_hint(&self) -> (usize, Option<usize>) {
+            (self.0, self.1)
+        }
+    }
+
+    // The two checks in each case are needed because the logic
+    // is different before the first call to `next()`.
+
+    let mut it = StubSizeHint(10, Some(10)).step_by(1);
+    assert_eq!(it.size_hint(), (10, Some(10)));
+    it.next();
+    assert_eq!(it.size_hint(), (9, Some(9)));
+
+    // exact multiple
+    let mut it = StubSizeHint(10, Some(10)).step_by(3);
+    assert_eq!(it.size_hint(), (4, Some(4)));
+    it.next();
+    assert_eq!(it.size_hint(), (3, Some(3)));
+
+    // larger base range, but not enough to get another element
+    let mut it = StubSizeHint(12, Some(12)).step_by(3);
+    assert_eq!(it.size_hint(), (4, Some(4)));
+    it.next();
+    assert_eq!(it.size_hint(), (3, Some(3)));
+
+    // smaller base range, so fewer resulting elements
+    let mut it = StubSizeHint(9, Some(9)).step_by(3);
+    assert_eq!(it.size_hint(), (3, Some(3)));
+    it.next();
+    assert_eq!(it.size_hint(), (2, Some(2)));
+
+    // infinite upper bound
+    let mut it = StubSizeHint(usize::MAX, None).step_by(1);
+    assert_eq!(it.size_hint(), (usize::MAX, None));
+    it.next();
+    assert_eq!(it.size_hint(), (usize::MAX-1, None));
+
+    // still infinite with larger step
+    let mut it = StubSizeHint(7, None).step_by(3);
+    assert_eq!(it.size_hint(), (3, None));
+    it.next();
+    assert_eq!(it.size_hint(), (2, None));
+
+    // propagates ExactSizeIterator
+    let a = [1,2,3,4,5];
+    let it = a.iter().step_by(2);
+    assert_eq!(it.len(), 3);
+
+    // Cannot be TrustedLen as a step greater than one makes an iterator
+    // with (usize::MAX, None) no longer meet the safety requirements
+    trait TrustedLenCheck { fn test(self) -> bool; }
+    impl<T:Iterator> TrustedLenCheck for T {
+        default fn test(self) -> bool { false }
+    }
+    impl<T:TrustedLen> TrustedLenCheck for T {
+        fn test(self) -> bool { true }
+    }
+    assert!(TrustedLenCheck::test(a.iter()));
+    assert!(!TrustedLenCheck::test(a.iter().step_by(1)));
+}
+
+#[test]
 fn test_filter_map() {
     let it = (0..).step_by(1).take(10)
         .filter_map(|x| if x % 2 == 0 { Some(x*x) } else { None });
diff --git a/src/libcore/tests/lib.rs b/src/libcore/tests/lib.rs
index c52155ead4f..d5d7d0f40b2 100644
--- a/src/libcore/tests/lib.rs
+++ b/src/libcore/tests/lib.rs
@@ -31,9 +31,11 @@
 #![feature(slice_patterns)]
 #![feature(sort_internals)]
 #![feature(sort_unstable)]
+#![feature(specialization)]
 #![feature(step_by)]
 #![feature(step_trait)]
 #![feature(test)]
+#![feature(trusted_len)]
 #![feature(try_from)]
 #![feature(unicode)]
 #![feature(unique)]