iter: introduce OptimisticCollect

Author: ljedrz

src/liballoc/collections/binary_heap.rs

@@ -146,7 +146,7 @@
 #![stable(feature = "rust1", since = "1.0.0")]
 
 use core::ops::{Deref, DerefMut};
-use core::iter::{FromIterator, FusedIterator};
+use core::iter::{FromIterator, FusedIterator, OptimisticCollect};
 use core::mem::{swap, size_of, ManuallyDrop};
 use core::ptr;
 use core::fmt;
@@ -1168,9 +1168,7 @@ impl<T: Ord> SpecExtend<BinaryHeap<T>> for BinaryHeap<T> {
 impl<T: Ord> BinaryHeap<T> {
     fn extend_desugared<I: IntoIterator<Item = T>>(&mut self, iter: I) {
         let iterator = iter.into_iter();
-        let (lower, _) = iterator.size_hint();
-
-        self.reserve(lower);
+        self.reserve(iterator.optimistic_collect_count());
 
         for elem in iterator {
             self.push(elem);

src/liballoc/collections/vec_deque.rs

@@ -9,7 +9,7 @@
 
 use core::cmp::Ordering;
 use core::fmt;
-use core::iter::{repeat_with, FromIterator, FusedIterator};
+use core::iter::{repeat_with, FromIterator, FusedIterator, OptimisticCollect};
 use core::mem;
 use core::ops::Bound::{Excluded, Included, Unbounded};
 use core::ops::{Index, IndexMut, RangeBounds, Try};
@@ -2597,8 +2597,7 @@ impl<A> IndexMut<usize> for VecDeque<A> {
 impl<A> FromIterator<A> for VecDeque<A> {
     fn from_iter<T: IntoIterator<Item = A>>(iter: T) -> VecDeque<A> {
         let iterator = iter.into_iter();
-        let (lower, _) = iterator.size_hint();
-        let mut deq = VecDeque::with_capacity(lower);
+        let mut deq = VecDeque::with_capacity(iterator.optimistic_collect_count());
         deq.extend(iterator);
         deq
     }

src/liballoc/lib.rs

@@ -114,6 +114,7 @@
 #![feature(maybe_uninit)]
 #![feature(alloc_layout_extra)]
 #![feature(try_trait)]
+#![feature(optimistic_collect)]
 
 // Allow testing this library
 

src/liballoc/string.rs

@@ -49,7 +49,7 @@
 use core::char::{decode_utf16, REPLACEMENT_CHARACTER};
 use core::fmt;
 use core::hash;
-use core::iter::{FromIterator, FusedIterator};
+use core::iter::{FromIterator, FusedIterator, OptimisticCollect};
 use core::ops::Bound::{Excluded, Included, Unbounded};
 use core::ops::{self, Add, AddAssign, Index, IndexMut, RangeBounds};
 use core::ptr;
@@ -1760,8 +1760,7 @@ impl<'a> FromIterator<Cow<'a, str>> for String {
 impl Extend<char> for String {
     fn extend<I: IntoIterator<Item = char>>(&mut self, iter: I) {
         let iterator = iter.into_iter();
-        let (lower_bound, _) = iterator.size_hint();
-        self.reserve(lower_bound);
+        self.reserve(iterator.optimistic_collect_count());
         iterator.for_each(move |c| self.push(c));
     }
 }

src/liballoc/vec.rs

@@ -60,7 +60,7 @@ use core::cmp::{self, Ordering};
 use core::fmt;
 use core::hash::{self, Hash};
 use core::intrinsics::{arith_offset, assume};
-use core::iter::{FromIterator, FusedIterator, TrustedLen};
+use core::iter::{FromIterator, FusedIterator, TrustedLen, OptimisticCollect};
 use core::marker::PhantomData;
 use core::mem;
 use core::ops::Bound::{Excluded, Included, Unbounded};
@@ -1813,8 +1813,8 @@ impl<T, I> SpecExtend<T, I> for Vec<T>
         let mut vector = match iterator.next() {
             None => return Vec::new(),
             Some(element) => {
-                let (lower, _) = iterator.size_hint();
-                let mut vector = Vec::with_capacity(lower.saturating_add(1));
+                let mut vector = Vec::with_capacity(
+                    iterator.optimistic_collect_count().saturating_add(1));
                 unsafe {
                     ptr::write(vector.get_unchecked_mut(0), element);
                     vector.set_len(1);
@@ -1933,8 +1933,7 @@ impl<T> Vec<T> {
         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));
+                self.reserve(iterator.optimistic_collect_count().saturating_add(1));
             }
             unsafe {
                 ptr::write(self.get_unchecked_mut(len), element);
@@ -2589,9 +2588,9 @@ impl<'a, I: Iterator> Drop for Splice<'a, I> {
 
             // There may be more elements. Use the lower bound as an estimate.
             // FIXME: Is the upper bound a better guess? Or something else?
-            let (lower_bound, _upper_bound) = self.replace_with.size_hint();
-            if lower_bound > 0  {
-                self.drain.move_tail(lower_bound);
+            let optimistic_collect_count = self.replace_with.optimistic_collect_count();
+            if optimistic_collect_count > 0  {
+                self.drain.move_tail(optimistic_collect_count);
                 if !self.drain.fill(&mut self.replace_with) {
                     return
                 }

src/libcore/benches/iter.rs

@@ -306,3 +306,8 @@ fn bench_skip_then_zip(b: &mut Bencher) {
         assert_eq!(s, 2009900);
     });
 }
+
+#[bench]
+fn bench_collect_to_result(b: &mut Bencher) {
+    b.iter(|| (0..100).map(|e| Ok(e)).collect::<Result<Vec<usize>, ()>>())
+}

src/libcore/iter/mod.rs

@@ -337,6 +337,8 @@ pub use self::traits::{ExactSizeIterator, Sum, Product};
 pub use self::traits::FusedIterator;
 #[unstable(feature = "trusted_len", issue = "37572")]
 pub use self::traits::TrustedLen;
+#[unstable(feature = "optimistic_collect", issue = "00000")]
+pub use self::traits::OptimisticCollect;
 
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::adapters::{Rev, Cycle, Chain, Zip, Map, Filter, FilterMap, Enumerate};

src/libcore/iter/traits/mod.rs

@@ -4,10 +4,12 @@ mod exact_size;
 mod collect;
 mod accum;
 mod marker;
+mod optimistic_collect;
 
 pub use self::iterator::Iterator;
 pub use self::double_ended::DoubleEndedIterator;
 pub use self::exact_size::ExactSizeIterator;
 pub use self::collect::{FromIterator, IntoIterator, Extend};
 pub use self::accum::{Sum, Product};
 pub use self::marker::{FusedIterator, TrustedLen};
+pub use self::optimistic_collect::OptimisticCollect;

src/libcore/iter/traits/optimistic_collect.rs

@@ -0,0 +1,22 @@
+/// A specialized trait designed to improve the estimates used when preallocating collections in
+/// cases where `size_hint` is too conservative. For instance, when collecting into an `Option` or a
+/// `Result`, the most common outcome is a non-empty collection, but the protocol allows `size_hint`
+/// to only provide a lower bound of `0`. `OptimisticCollect` can be specialized for such cases in
+/// order to optimize the creation of the resulting collections without breaking `Iterator` rules.
+#[unstable(feature = "optimistic_collect", issue = "00000")]
+pub trait OptimisticCollect: Iterator {
+    /// Provides an estimate of the size of the iterator for the purposes of preallocating
+    /// collections that can be built from it. By default it provides the lower bound of
+    /// `size_hint`.
+    fn optimistic_collect_count(&self) -> usize;
+}
+
+#[unstable(feature = "optimistic_collect", issue = "00000")]
+impl<I: Iterator> OptimisticCollect for I {
+    default fn optimistic_collect_count(&self) -> usize { self.size_hint().0 }
+}
+
+#[unstable(feature = "optimistic_collect", issue = "00000")]
+impl<I: Iterator> OptimisticCollect for &mut I {
+    default fn optimistic_collect_count(&self) -> usize { (**self).size_hint().0 }
+}

src/libcore/option.rs

@@ -135,7 +135,7 @@
 
 #![stable(feature = "rust1", since = "1.0.0")]
 
-use iter::{FromIterator, FusedIterator, TrustedLen};
+use iter::{FromIterator, FusedIterator, TrustedLen, OptimisticCollect};
 use {hint, mem, ops::{self, Deref}};
 use pin::Pin;
 
@@ -1356,6 +1356,17 @@ impl<A, V: FromIterator<A>> FromIterator<Option<A>> for Option<V> {
             }
         }
 
+        impl<T, Iter: Iterator<Item = Option<T>>> OptimisticCollect for Adapter<Iter> {
+            #[inline]
+            fn optimistic_collect_count(&self) -> usize {
+                if self.found_none {
+                    0
+                } else {
+                    self.iter.optimistic_collect_count()
+                }
+            }
+        }
+
         let mut adapter = Adapter { iter: iter.into_iter(), found_none: false };
         let v: V = FromIterator::from_iter(adapter.by_ref());
 

src/libcore/result.rs

@@ -231,7 +231,7 @@
 #![stable(feature = "rust1", since = "1.0.0")]
 
 use fmt;
-use iter::{FromIterator, FusedIterator, TrustedLen};
+use iter::{FromIterator, FusedIterator, TrustedLen, OptimisticCollect};
 use ops::{self, Deref};
 
 /// `Result` is a type that represents either success ([`Ok`]) or failure ([`Err`]).
@@ -1233,6 +1233,17 @@ impl<A, E, V: FromIterator<A>> FromIterator<Result<A, E>> for Result<V, E> {
             }
         }
 
+        impl<T, E, Iter: Iterator<Item = Result<T, E>>> OptimisticCollect for Adapter<Iter, E> {
+            #[inline]
+            fn optimistic_collect_count(&self) -> usize {
+                if self.err.is_some() {
+                    0
+                } else {
+                    self.iter.optimistic_collect_count()
+                }
+            }
+        }
+
         let mut adapter = Adapter { iter: iter.into_iter(), err: None };
         let v: V = FromIterator::from_iter(adapter.by_ref());
 

src/libstd/lib.rs

@@ -297,6 +297,7 @@
 #![feature(non_exhaustive)]
 #![feature(alloc_layout_extra)]
 #![feature(maybe_uninit)]
+#![feature(optimistic_collect)]
 #![cfg_attr(all(target_vendor = "fortanix", target_env = "sgx"),
             feature(global_asm, range_contains, slice_index_methods,
                     decl_macro, coerce_unsized, sgx_platform, ptr_wrapping_offset_from))]

src/libstd/sys_common/wtf8.rs

@@ -21,7 +21,7 @@ use borrow::Cow;
 use char;
 use fmt;
 use hash::{Hash, Hasher};
-use iter::FromIterator;
+use iter::{FromIterator, OptimisticCollect};
 use mem;
 use ops;
 use rc::Rc;
@@ -385,9 +385,8 @@ impl FromIterator<CodePoint> for Wtf8Buf {
 impl Extend<CodePoint> for Wtf8Buf {
     fn extend<T: IntoIterator<Item=CodePoint>>(&mut self, iter: T) {
         let iterator = iter.into_iter();
-        let (low, _high) = iterator.size_hint();
         // Lower bound of one byte per code point (ASCII only)
-        self.bytes.reserve(low);
+        self.bytes.reserve(iterator.optimistic_collect_count());
         for code_point in iterator {
             self.push(code_point);
         }