Added set_size_hint method

Author: peterjoel

src/adaptors/mod.rs

@@ -1276,3 +1276,49 @@ where
         }
     }
 }
+
+/// An iterator adaptor that applies the given `size_hint` to an iterator.
+///
+/// See [`.set_size_hint()`](../trait.Itertools.html#method.set_size_hint.html)
+/// for more information.
+#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
+pub struct SetSizeHint<I> {
+    iter: I,
+    size_hint: (usize, Option<usize>),
+}
+
+/// Create a new `SetSizeHint` iterator adapter
+pub fn set_size_hint<I>(iter: I, min: usize, max: Option<usize>) -> SetSizeHint<I> {
+    SetSizeHint { iter, size_hint: (min, max) }
+}
+
+impl<I> Iterator for SetSizeHint<I>
+where
+    I: Iterator,
+{
+    type Item = I::Item;
+
+    #[inline(always)]
+    fn next(&mut self) -> Option<Self::Item> {
+        self.iter.next()
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.size_hint
+    }
+}
+
+impl<I> ExactSizeIterator for SetSizeHint<I>
+where
+    I: ExactSizeIterator,
+{ }
+
+impl<I> DoubleEndedIterator for SetSizeHint<I>
+where
+    I: DoubleEndedIterator,
+{
+    #[inline(always)]
+    fn next_back(&mut self) -> Option<Self::Item> {
+        self.iter.next_back()
+    }
+}

src/lib.rs

@@ -78,22 +78,23 @@ pub use std::iter as __std_iter;
 /// The concrete iterator types.
 pub mod structs {
     pub use adaptors::{
+        Batching,
+        Coalesce,
         Dedup,
         Interleave,
         InterleaveShortest,
-        Product,
-        PutBack,
-        Batching,
         MapInto,
         MapResults,
         Merge,
         MergeBy,
+        Positions,
+        Product,
+        PutBack,
+        SetSizeHint,
         TakeWhileRef,
-        WhileSome,
-        Coalesce,
         TupleCombinations,
-        Positions,
         Update,
+        WhileSome,
     };
     #[allow(deprecated)]
     pub use adaptors::Step;
@@ -1902,6 +1903,40 @@ pub trait Itertools : Iterator {
         v.into_iter()
     }
 
+    /// Explicitly provide a size hint for an iterator.
+    ///
+    /// This can be useful in situations where the size hint cannot be deduced
+    /// automatically, but where the programmer knows the bounds on the number
+    /// of elements.
+    ///
+    /// ```
+    /// use itertools::{self, Itertools};
+    ///
+    /// let data = vec![1, 2, 6, 7, 2];
+    ///
+    /// let result: Vec<i32> = data.iter()
+    ///     // The `FlatMap` adapter is not able to deduce the size hint itself
+    ///     // so `size_hint()` would return `(0, None)`.
+    ///     .flat_map(|&x| {
+    ///         let repeats = if x % 2 == 0 { 1 } else { 3 };
+    ///         itertools::repeat_n(x, repeats)
+    ///     })
+    ///     // But we know the bounds on the max and min number of items in the
+    ///     // resulting iterator, so we can provide that information:
+    ///     .set_size_hint(data.len(), Some(3 * data.len()))
+    ///     // The `Vec` should not excessively re-allocate, while collecting
+    ///     // since it now knows the minimum and maximum number of items.
+    ///     .collect();
+    ///
+    /// assert_eq!(result, vec![1, 1, 1, 2, 6, 7, 7, 7, 2]);
+    /// ```
+    ///
+    fn set_size_hint(self, min: usize, max: Option<usize>) -> SetSizeHint<Self>
+        where Self: Sized,
+    {
+        adaptors::set_size_hint(self, min, max)
+    }
+
     /// Collect all iterator elements into one of two
     /// partitions. Unlike `Iterator::partition`, each partition may
     /// have a distinct type.