Rollup merge of rust-lang#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 rust-lang#47082

r? @alexcrichton

Author: kennytm

src/libcore/iter/mod.rs

@@ -2286,16 +2286,7 @@ impl<I> Iterator for Take<I> where I: Iterator{
 
     #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
-        let (lower, upper) = self.iter.size_hint();
-
-        let lower = cmp::min(lower, self.n);
-
-        let upper = match upper {
-            Some(x) if x < self.n => Some(x),
-            _ => Some(self.n)
-        };
-
-        (lower, upper)
+        TakeImpl::size_hint(self)
     }
 
     #[inline]
@@ -2316,12 +2307,50 @@ impl<I> Iterator for Take<I> where I: Iterator{
     }
 }
 
+trait TakeImpl {
+    fn size_hint(&self) -> (usize, Option<usize>);
+}
+
+impl<I: Iterator> TakeImpl for Take<I> {
+    #[inline]
+    default fn size_hint(&self) -> (usize, Option<usize>) {
+        let (lower, upper) = self.iter.size_hint();
+
+        let lower = cmp::min(lower, self.n);
+
+        let upper = match upper {
+            Some(x) if x < self.n => Some(x),
+            _ => Some(self.n)
+        };
+
+        (lower, upper)
+    }
+}
+
+impl<I: TrustedLen> TakeImpl for Take<I> {
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let (lower, upper) = self.iter.size_hint();
+        match upper {
+            None => (self.n, Some(self.n)),
+            Some(x) => {
+                debug_assert_eq!(x, lower);
+                let count = cmp::min(lower, self.n);
+                (count, Some(count))
+            }
+        }
+    }
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 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

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;

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

src/libcore/iter/traits.rs

@@ -972,7 +972,8 @@ impl<'a, I: FusedIterator + ?Sized> FusedIterator for &'a mut I {}
 /// The upper bound must only be [`None`] if the actual iterator length is
 /// larger than [`usize::MAX`].
 ///
-/// 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
 ///

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]
@@ -1465,6 +1488,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]

src/test/codegen/repeat-trusted-len.rs

@@ -0,0 +1,31 @@
+// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// compile-flags: -O
+// ignore-tidy-linelength
+
+#![crate_type = "lib"]
+
+use std::iter;
+
+// CHECK-LABEL: @repeat_take_collect
+#[no_mangle]
+pub fn repeat_take_collect() -> Vec<u8> {
+// CHECK: call void @llvm.memset.p0i8
+    iter::repeat(42).take(100000).collect()
+}
+
+// CHECK-LABEL: @range_from_take_collect
+#[no_mangle]
+pub fn range_from_take_collect() -> Vec<u8> {
+// CHECK: %[[SPLATINSERT:.*]] = insertelement <{{[0-9]+}} x i8> undef, i8 %{{.*}}, i32 0
+// CHECK-NEXT: %{{.*}} = shufflevector <[[WIDTH:[0-9]+]] x i8> %[[SPLATINSERT]], <[[WIDTH]] x i8> undef, <[[WIDTH]] x i32> zeroinitializer
+    (0..).take(100000).collect()
+}