diff --git a/src/adaptors/mod.rs b/src/adaptors/mod.rs index 59aebfaf0..0b5845623 100644 --- a/src/adaptors/mod.rs +++ b/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 { + iter: I, + size_hint: (usize, Option), +} + +/// Create a new `SetSizeHint` iterator adapter +pub fn set_size_hint(iter: I, min: usize, max: Option) -> SetSizeHint { + SetSizeHint { iter, size_hint: (min, max) } +} + +impl Iterator for SetSizeHint +where + I: Iterator, +{ + type Item = I::Item; + + #[inline(always)] + fn next(&mut self) -> Option { + self.iter.next() + } + + fn size_hint(&self) -> (usize, Option) { + self.size_hint + } +} + +impl ExactSizeIterator for SetSizeHint +where + I: ExactSizeIterator, +{ } + +impl DoubleEndedIterator for SetSizeHint +where + I: DoubleEndedIterator, +{ + #[inline(always)] + fn next_back(&mut self) -> Option { + self.iter.next_back() + } +} diff --git a/src/lib.rs b/src/lib.rs index aa01a4f91..72515c0e0 100644 --- a/src/lib.rs +++ b/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 = 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) -> SetSizeHint + 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.