Auto merge of #42167 - scottmcm:iter-stepby-sizehint, r=alexcrichton

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`.

A small non-breaking subset of https://github.com/rust-lang/rust/pull/42110 (which I closed).

Includes two small documentation changes @ivandardi requested on that PR.

r? @alexcrichton

4 files changed, 96 insertions, 1 deletions

diff --git a/src/doc/unstable-book/src/SUMMARY.md b/src/doc/unstable-book/src/SUMMARY.md
index f944675e5eb..8b70e8be38a 100644
--- a/src/doc/unstable-book/src/SUMMARY.md
+++ b/src/doc/unstable-book/src/SUMMARY.md
@@ -153,6 +153,7 @@
     - [io](library-features/io.md)
     - [ip](library-features/ip.md)
     - [iter_rfind](library-features/iter-rfind.md)
+    - [iterator_step_by](library-features/iterator-step-by.md)
     - [libstd_io_internals](library-features/libstd-io-internals.md)
     - [libstd_sys_internals](library-features/libstd-sys-internals.md)
     - [libstd_thread_internals](library-features/libstd-thread-internals.md)
diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs
index 5eefa59e7ea..07aed65f7a0 100644
--- a/src/libcore/iter/mod.rs
+++ b/src/libcore/iter/mod.rs
@@ -520,7 +520,7 @@ impl<I> Iterator for Cycle<I> where I: Clone + Iterator {
 #[unstable(feature = "fused", issue = "35602")]
 impl<I> FusedIterator for Cycle<I> where I: Clone + Iterator {}
 
-/// An iterator that steps by n elements every iteration.
+/// An adapter for stepping iterators by a custom amount.
 ///
 /// This `struct` is created by the [`step_by`] method on [`Iterator`]. See
 /// its documentation for more.
@@ -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 8c4cd1d0c84..e9f62dfbaed 100644
--- a/src/libcore/tests/lib.rs
+++ b/src/libcore/tests/lib.rs
@@ -32,9 +32,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)]