Rollup merge of #47944 - oberien:unboundediterator-trustedlen, r=bluss

Implement TrustedLen for Take<Repeat> and Take<RangeFrom>

This will allow optimization of simple `repeat(x).take(n).collect()` iterators, which are currently not vectorized and have capacity checks.

This will only support a few aggregates on `Repeat` and `RangeFrom`, which might be enough for simple cases, but doesn't optimize more complex ones. Namely, Cycle, StepBy, Filter, FilterMap, Peekable, SkipWhile, Skip, FlatMap, Fuse and Inspect are not marked `TrustedLen` when the inner iterator is infinite.

Previous discussion can be found in #47082

r? @alexcrichton

5 files changed, 56 insertions, 2 deletions

diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs
index 7314fac282b..bf8367d85fd 100644
--- a/src/libcore/iter/mod.rs
+++ b/src/libcore/iter/mod.rs
@@ -2322,6 +2322,9 @@ impl<I> ExactSizeIterator for Take<I> where I: ExactSizeIterator {}
 #[unstable(feature = "fused", issue = "35602")]
 impl<I> FusedIterator for Take<I> where I: FusedIterator {}
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<I: TrustedLen> TrustedLen for Take<I> {}
+
 /// An iterator to maintain state while iterating another iterator.
 ///
 /// This `struct` is created by the [`scan`] method on [`Iterator`]. See its
diff --git a/src/libcore/iter/range.rs b/src/libcore/iter/range.rs
index 3b034efcce1..5b4dc9b2c9a 100644
--- a/src/libcore/iter/range.rs
+++ b/src/libcore/iter/range.rs
@@ -325,6 +325,9 @@ impl<A: Step> Iterator for ops::RangeFrom<A> {
 #[unstable(feature = "fused", issue = "35602")]
 impl<A: Step> FusedIterator for ops::RangeFrom<A> {}
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<A: Step> TrustedLen for ops::RangeFrom<A> {}
+
 #[unstable(feature = "inclusive_range", reason = "recently added, follows RFC", issue = "28237")]
 impl<A: Step> Iterator for ops::RangeInclusive<A> {
     type Item = A;
diff --git a/src/libcore/iter/sources.rs b/src/libcore/iter/sources.rs
index b405f35d5e4..b05a893e661 100644
--- a/src/libcore/iter/sources.rs
+++ b/src/libcore/iter/sources.rs
@@ -44,6 +44,9 @@ impl<A: Clone> DoubleEndedIterator for Repeat<A> {
 #[unstable(feature = "fused", issue = "35602")]
 impl<A: Clone> FusedIterator for Repeat<A> {}
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<A: Clone> TrustedLen for Repeat<A> {}
+
 /// Creates a new iterator that endlessly repeats a single element.
 ///
 /// The `repeat()` function repeats a single value over and over and over and
diff --git a/src/libcore/iter/traits.rs b/src/libcore/iter/traits.rs
index 11e668d228c..be4889f2487 100644
--- a/src/libcore/iter/traits.rs
+++ b/src/libcore/iter/traits.rs
@@ -970,9 +970,11 @@ impl<'a, I: FusedIterator + ?Sized> FusedIterator for &'a mut I {}
 /// The iterator reports a size hint where it is either exact
 /// (lower bound is equal to upper bound), or the upper bound is [`None`].
 /// The upper bound must only be [`None`] if the actual iterator length is
-/// larger than [`usize::MAX`].
+/// larger than [`usize::MAX`]. In that case, the lower bound must be
+/// [`usize::MAX`], resulting in a [`.size_hint`] of `(usize::MAX, None)`.
 ///
-/// The iterator must produce exactly the number of elements it reported.
+/// The iterator must produce exactly the number of elements it reported
+/// or diverge before reaching the end.
 ///
 /// # Safety
 ///
diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs
index dc866d180bf..0373c8922b2 100644
--- a/src/libcore/tests/iter.rs
+++ b/src/libcore/tests/iter.rs
@@ -1371,6 +1371,29 @@ fn test_range_from_nth() {
     assert_eq!(r, 16..);
     assert_eq!(r.nth(10), Some(26));
     assert_eq!(r, 27..);
+
+    assert_eq!((0..).size_hint(), (usize::MAX, None));
+}
+
+fn is_trusted_len<I: TrustedLen>(_: I) {}
+
+#[test]
+fn test_range_from_take() {
+    let mut it = (0..).take(3);
+    assert_eq!(it.next(), Some(0));
+    assert_eq!(it.next(), Some(1));
+    assert_eq!(it.next(), Some(2));
+    assert_eq!(it.next(), None);
+    is_trusted_len((0..).take(3));
+    assert_eq!((0..).take(3).size_hint(), (3, Some(3)));
+    assert_eq!((0..).take(0).size_hint(), (0, Some(0)));
+    assert_eq!((0..).take(usize::MAX).size_hint(), (usize::MAX, Some(usize::MAX)));
+}
+
+#[test]
+fn test_range_from_take_collect() {
+    let v: Vec<_> = (0..).take(3).collect();
+    assert_eq!(v, vec![0, 1, 2]);
 }
 
 #[test]
@@ -1485,6 +1508,26 @@ fn test_repeat() {
     assert_eq!(it.next(), Some(42));
     assert_eq!(it.next(), Some(42));
     assert_eq!(it.next(), Some(42));
+    assert_eq!(repeat(42).size_hint(), (usize::MAX, None));
+}
+
+#[test]
+fn test_repeat_take() {
+    let mut it = repeat(42).take(3);
+    assert_eq!(it.next(), Some(42));
+    assert_eq!(it.next(), Some(42));
+    assert_eq!(it.next(), Some(42));
+    assert_eq!(it.next(), None);
+    is_trusted_len(repeat(42).take(3));
+    assert_eq!(repeat(42).take(3).size_hint(), (3, Some(3)));
+    assert_eq!(repeat(42).take(0).size_hint(), (0, Some(0)));
+    assert_eq!(repeat(42).take(usize::MAX).size_hint(), (usize::MAX, Some(usize::MAX)));
+}
+
+#[test]
+fn test_repeat_take_collect() {
+    let v: Vec<_> = repeat(42).take(3).collect();
+    assert_eq!(v, vec![42, 42, 42]);
 }
 
 #[test]