Auto merge of #37306 - bluss:trusted-len, r=alexcrichton

Add Iterator trait TrustedLen to enable better FromIterator / Extend

This trait attempts to improve FromIterator / Extend code by enabling it to trust the iterator to produce an exact number of elements, which means that reallocation needs to happen only once and is moved out of the loop.

`TrustedLen` differs from `ExactSizeIterator` in that it attempts to include _more_ iterators by allowing for the case that the iterator's len does not fit in `usize`. Consumers must check for this case (for example they could panic, since they can't allocate a collection of that size).

For example, chain can be TrustedLen and all numerical ranges can be TrustedLen. All they need to do is to report an exact size if it fits in `usize`, and `None` as the upper bound otherwise.

The trait describes its contract like this:

```
An iterator that reports an accurate length using size_hint.

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

The iterator must produce exactly the number of elements it reported.

This trait must only be implemented when the contract is upheld.
Consumers of this trait must inspect `.size_hint()`’s upper bound.
```

Fixes #37232

Author: bors

src/libcollections/lib.rs

@@ -50,6 +50,7 @@
 #![feature(specialization)]
 #![feature(staged_api)]
 #![feature(step_by)]
+#![feature(trusted_len)]
 #![feature(unicode)]
 #![feature(unique)]
 #![cfg_attr(test, feature(rand, test))]

src/libcollections/vec.rs

@@ -75,15 +75,14 @@ use core::cmp::Ordering;
 use core::fmt;
 use core::hash::{self, Hash};
 use core::intrinsics::{arith_offset, assume};
-use core::iter::{FromIterator, FusedIterator};
+use core::iter::{FromIterator, FusedIterator, TrustedLen};
 use core::mem;
 use core::ops::{Index, IndexMut};
 use core::ops;
 use core::ptr;
 use core::ptr::Shared;
 use core::slice;
 
-use super::SpecExtend;
 use super::range::RangeArgument;
 
 /// A contiguous growable array type, written `Vec<T>` but pronounced 'vector'.
@@ -1245,26 +1244,7 @@ impl<T: Clone> Vec<T> {
     /// ```
     #[stable(feature = "vec_extend_from_slice", since = "1.6.0")]
     pub fn extend_from_slice(&mut self, other: &[T]) {
-        self.reserve(other.len());
-
-        // Unsafe code so this can be optimised to a memcpy (or something
-        // similarly fast) when T is Copy. LLVM is easily confused, so any
-        // extra operations during the loop can prevent this optimisation.
-        unsafe {
-            let len = self.len();
-            let ptr = self.get_unchecked_mut(len) as *mut T;
-            // Use SetLenOnDrop to work around bug where compiler
-            // may not realize the store through `ptr` trough self.set_len()
-            // don't alias.
-            let mut local_len = SetLenOnDrop::new(&mut self.len);
-
-            for i in 0..other.len() {
-                ptr::write(ptr.offset(i as isize), other.get_unchecked(i).clone());
-                local_len.increment_len(1);
-            }
-
-            // len set by scope guard
-        }
+        self.extend(other.iter().cloned())
     }
 }
 
@@ -1606,19 +1586,25 @@ impl<'a, T> IntoIterator for &'a mut Vec<T> {
 impl<T> Extend<T> for Vec<T> {
     #[inline]
     fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
-        <Self as SpecExtend<I>>::spec_extend(self, iter);
+        self.extend_desugared(iter.into_iter())
     }
 }
 
-impl<I: IntoIterator> SpecExtend<I> for Vec<I::Item> {
-    default fn spec_extend(&mut self, iter: I) {
-        self.extend_desugared(iter.into_iter())
-    }
+trait IsTrustedLen : Iterator {
+    fn trusted_len(&self) -> Option<usize> { None }
 }
+impl<I> IsTrustedLen for I where I: Iterator { }
 
-impl<T> SpecExtend<Vec<T>> for Vec<T> {
-    fn spec_extend(&mut self, ref mut other: Vec<T>) {
-        self.append(other);
+impl<I> IsTrustedLen for I where I: TrustedLen
+{
+    fn trusted_len(&self) -> Option<usize> {
+        let (low, high) = self.size_hint();
+        if let Some(high_value) = high {
+            debug_assert_eq!(low, high_value,
+                             "TrustedLen iterator's size hint is not exact: {:?}",
+                             (low, high));
+        }
+        high
     }
 }
 
@@ -1629,16 +1615,30 @@ impl<T> Vec<T> {
         //      for item in iterator {
         //          self.push(item);
         //      }
-        while let Some(element) = iterator.next() {
-            let len = self.len();
-            if len == self.capacity() {
-                let (lower, _) = iterator.size_hint();
-                self.reserve(lower.saturating_add(1));
-            }
+        if let Some(additional) = iterator.trusted_len() {
+            self.reserve(additional);
             unsafe {
-                ptr::write(self.get_unchecked_mut(len), element);
-                // NB can't overflow since we would have had to alloc the address space
-                self.set_len(len + 1);
+                let mut ptr = self.as_mut_ptr().offset(self.len() as isize);
+                let mut local_len = SetLenOnDrop::new(&mut self.len);
+                for element in iterator {
+                    ptr::write(ptr, element);
+                    ptr = ptr.offset(1);
+                    // NB can't overflow since we would have had to alloc the address space
+                    local_len.increment_len(1);
+                }
+            }
+        } else {
+            while let Some(element) = iterator.next() {
+                let len = self.len();
+                if len == self.capacity() {
+                    let (lower, _) = iterator.size_hint();
+                    self.reserve(lower.saturating_add(1));
+                }
+                unsafe {
+                    ptr::write(self.get_unchecked_mut(len), element);
+                    // NB can't overflow since we would have had to alloc the address space
+                    self.set_len(len + 1);
+                }
             }
         }
     }
@@ -1647,24 +1647,7 @@ impl<T> Vec<T> {
 #[stable(feature = "extend_ref", since = "1.2.0")]
 impl<'a, T: 'a + Copy> Extend<&'a T> for Vec<T> {
     fn extend<I: IntoIterator<Item = &'a T>>(&mut self, iter: I) {
-        <I as SpecExtendVec<T>>::extend_vec(iter, self);
-    }
-}
-
-// helper trait for specialization of Vec's Extend impl
-trait SpecExtendVec<T> {
-    fn extend_vec(self, vec: &mut Vec<T>);
-}
-
-impl <'a, T: 'a + Copy, I: IntoIterator<Item=&'a T>> SpecExtendVec<T> for I {
-    default fn extend_vec(self, vec: &mut Vec<T>) {
-        vec.extend(self.into_iter().cloned());
-    }
-}
-
-impl<'a, T: Copy> SpecExtendVec<T> for &'a [T] {
-    fn extend_vec(self, vec: &mut Vec<T>) {
-        vec.extend_from_slice(self);
+        self.extend(iter.into_iter().map(|&x| x))
     }
 }
 
@@ -1988,6 +1971,9 @@ impl<T> ExactSizeIterator for IntoIter<T> {}
 #[unstable(feature = "fused", issue = "35602")]
 impl<T> FusedIterator for IntoIter<T> {}
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<T> TrustedLen for IntoIter<T> {}
+
 #[stable(feature = "vec_into_iter_clone", since = "1.8.0")]
 impl<T: Clone> Clone for IntoIter<T> {
     fn clone(&self) -> IntoIter<T> {

src/libcore/iter/mod.rs

@@ -328,6 +328,8 @@ pub use self::traits::{FromIterator, IntoIterator, DoubleEndedIterator, Extend};
 pub use self::traits::{ExactSizeIterator, Sum, Product};
 #[unstable(feature = "fused", issue = "35602")]
 pub use self::traits::FusedIterator;
+#[unstable(feature = "trusted_len", issue = "37572")]
+pub use self::traits::TrustedLen;
 
 mod iterator;
 mod range;
@@ -372,6 +374,10 @@ impl<I> ExactSizeIterator for Rev<I>
 impl<I> FusedIterator for Rev<I>
     where I: FusedIterator + DoubleEndedIterator {}
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<I> TrustedLen for Rev<I>
+    where I: TrustedLen + DoubleEndedIterator {}
+
 /// An iterator that clones the elements of an underlying iterator.
 ///
 /// This `struct` is created by the [`cloned()`] method on [`Iterator`]. See its
@@ -438,6 +444,12 @@ unsafe impl<'a, I, T: 'a> TrustedRandomAccess for Cloned<I>
     fn may_have_side_effect() -> bool { true }
 }
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<'a, I, T: 'a> TrustedLen for Cloned<I>
+    where I: TrustedLen<Item=&'a T>,
+          T: Clone
+{}
+
 /// An iterator that repeats endlessly.
 ///
 /// This `struct` is created by the [`cycle()`] method on [`Iterator`]. See its
@@ -667,6 +679,11 @@ impl<A, B> FusedIterator for Chain<A, B>
           B: FusedIterator<Item=A::Item>,
 {}
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<A, B> TrustedLen for Chain<A, B>
+    where A: TrustedLen, B: TrustedLen<Item=A::Item>,
+{}
+
 /// An iterator that iterates two other iterators simultaneously.
 ///
 /// This `struct` is created by the [`zip()`] method on [`Iterator`]. See its
@@ -884,6 +901,11 @@ unsafe impl<A, B> TrustedRandomAccess for Zip<A, B>
 impl<A, B> FusedIterator for Zip<A, B>
     where A: FusedIterator, B: FusedIterator, {}
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<A, B> TrustedLen for Zip<A, B>
+    where A: TrustedLen, B: TrustedLen,
+{}
+
 /// An iterator that maps the values of `iter` with `f`.
 ///
 /// This `struct` is created by the [`map()`] method on [`Iterator`]. See its
@@ -991,6 +1013,11 @@ impl<B, I: ExactSizeIterator, F> ExactSizeIterator for Map<I, F>
 impl<B, I: FusedIterator, F> FusedIterator for Map<I, F>
     where F: FnMut(I::Item) -> B {}
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<B, I, F> TrustedLen for Map<I, F>
+    where I: TrustedLen,
+          F: FnMut(I::Item) -> B {}
+
 #[doc(hidden)]
 unsafe impl<B, I, F> TrustedRandomAccess for Map<I, F>
     where I: TrustedRandomAccess,
@@ -1227,6 +1254,12 @@ unsafe impl<I> TrustedRandomAccess for Enumerate<I>
 #[unstable(feature = "fused", issue = "35602")]
 impl<I> FusedIterator for Enumerate<I> where I: FusedIterator {}
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<I> TrustedLen for Enumerate<I>
+    where I: TrustedLen,
+{}
+
+
 /// An iterator with a `peek()` that returns an optional reference to the next
 /// element.
 ///

src/libcore/iter/range.rs

@@ -12,7 +12,7 @@ use mem;
 use ops::{self, Add, Sub};
 use usize;
 
-use super::FusedIterator;
+use super::{FusedIterator, TrustedLen};
 
 /// Objects that can be stepped over in both directions.
 ///
@@ -480,6 +480,22 @@ macro_rules! range_incl_exact_iter_impl {
     )*)
 }
 
+macro_rules! range_trusted_len_impl {
+    ($($t:ty)*) => ($(
+        #[unstable(feature = "trusted_len", issue = "37572")]
+        unsafe impl TrustedLen for ops::Range<$t> { }
+    )*)
+}
+
+macro_rules! range_incl_trusted_len_impl {
+    ($($t:ty)*) => ($(
+        #[unstable(feature = "inclusive_range",
+                   reason = "recently added, follows RFC",
+                   issue = "28237")]
+        unsafe impl TrustedLen for ops::RangeInclusive<$t> { }
+    )*)
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<A: Step> Iterator for ops::Range<A> where
     for<'a> &'a A: Add<&'a A, Output = A>
@@ -513,6 +529,13 @@ impl<A: Step> Iterator for ops::Range<A> where
 range_exact_iter_impl!(usize u8 u16 u32 isize i8 i16 i32);
 range_incl_exact_iter_impl!(u8 u16 i8 i16);
 
+// These macros generate `TrustedLen` impls.
+//
+// They need to guarantee that .size_hint() is either exact, or that
+// the upper bound is None when it does not fit the type limits.
+range_trusted_len_impl!(usize isize u8 i8 u16 i16 u32 i32 i64 u64);
+range_incl_trusted_len_impl!(usize isize u8 i8 u16 i16 u32 i32 i64 u64);
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<A: Step + Clone> DoubleEndedIterator for ops::Range<A> where
     for<'a> &'a A: Add<&'a A, Output = A>,

src/libcore/iter/traits.rs

@@ -665,3 +665,22 @@ pub trait FusedIterator: Iterator {}
 
 #[unstable(feature = "fused", issue = "35602")]
 impl<'a, I: FusedIterator + ?Sized> FusedIterator for &'a mut I {}
+
+/// An iterator that reports an accurate length using size_hint.
+///
+/// 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`.
+///
+/// The iterator must produce exactly the number of elements it reported.
+///
+/// # Safety
+///
+/// This trait must only be implemented when the contract is upheld.
+/// Consumers of this trait must inspect `.size_hint()`’s upper bound.
+#[unstable(feature = "trusted_len", issue = "37572")]
+pub unsafe trait TrustedLen : Iterator {}
+
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<'a, I: TrustedLen + ?Sized> TrustedLen for &'a mut I {}

src/libcore/option.rs

@@ -145,7 +145,7 @@
 
 #![stable(feature = "rust1", since = "1.0.0")]
 
-use iter::{FromIterator, FusedIterator};
+use iter::{FromIterator, FusedIterator, TrustedLen};
 use mem;
 
 // Note that this is not a lang item per se, but it has a hidden dependency on
@@ -803,6 +803,7 @@ impl<A> DoubleEndedIterator for Item<A> {
 
 impl<A> ExactSizeIterator for Item<A> {}
 impl<A> FusedIterator for Item<A> {}
+unsafe impl<A> TrustedLen for Item<A> {}
 
 /// An iterator over a reference of the contained item in an [`Option`].
 ///
@@ -833,6 +834,9 @@ impl<'a, A> ExactSizeIterator for Iter<'a, A> {}
 #[unstable(feature = "fused", issue = "35602")]
 impl<'a, A> FusedIterator for Iter<'a, A> {}
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<'a, A> TrustedLen for Iter<'a, A> {}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<'a, A> Clone for Iter<'a, A> {
     fn clone(&self) -> Iter<'a, A> {
@@ -868,6 +872,8 @@ impl<'a, A> ExactSizeIterator for IterMut<'a, A> {}
 
 #[unstable(feature = "fused", issue = "35602")]
 impl<'a, A> FusedIterator for IterMut<'a, A> {}
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<'a, A> TrustedLen for IterMut<'a, A> {}
 
 /// An iterator over the item contained inside an [`Option`].
 ///
@@ -898,6 +904,9 @@ impl<A> ExactSizeIterator for IntoIter<A> {}
 #[unstable(feature = "fused", issue = "35602")]
 impl<A> FusedIterator for IntoIter<A> {}
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<A> TrustedLen for IntoIter<A> {}
+
 /////////////////////////////////////////////////////////////////////////////
 // FromIterator
 /////////////////////////////////////////////////////////////////////////////

src/libcore/result.rs

@@ -249,7 +249,7 @@
 #![stable(feature = "rust1", since = "1.0.0")]
 
 use fmt;
-use iter::{FromIterator, FusedIterator};
+use iter::{FromIterator, FusedIterator, TrustedLen};
 
 /// `Result` is a type that represents either success (`Ok`) or failure (`Err`).
 ///
@@ -924,6 +924,9 @@ impl<'a, T> ExactSizeIterator for Iter<'a, T> {}
 #[unstable(feature = "fused", issue = "35602")]
 impl<'a, T> FusedIterator for Iter<'a, T> {}
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<'a, A> TrustedLen for Iter<'a, A> {}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<'a, T> Clone for Iter<'a, T> {
     fn clone(&self) -> Iter<'a, T> { Iter { inner: self.inner } }
@@ -962,6 +965,9 @@ impl<'a, T> ExactSizeIterator for IterMut<'a, T> {}
 #[unstable(feature = "fused", issue = "35602")]
 impl<'a, T> FusedIterator for IterMut<'a, T> {}
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<'a, A> TrustedLen for IterMut<'a, A> {}
+
 /// An iterator over the value in a [`Ok`] variant of a [`Result`]. This struct is
 /// created by the [`into_iter`] method on [`Result`][`Result`] (provided by
 /// the [`IntoIterator`] trait).
@@ -999,6 +1005,9 @@ impl<T> ExactSizeIterator for IntoIter<T> {}
 #[unstable(feature = "fused", issue = "35602")]
 impl<T> FusedIterator for IntoIter<T> {}
 
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<A> TrustedLen for IntoIter<A> {}
+
 /////////////////////////////////////////////////////////////////////////////
 // FromIterator
 /////////////////////////////////////////////////////////////////////////////