Add an Iterator::batch transformer and explicit size hints for FromIterator and Extend

src/libcollections/dlist.rs

@@ -582,15 +582,15 @@ impl<A> DoubleEndedIterator<A> for MoveItems<A> {
 }
 
 impl<A> FromIterator<A> for DList<A> {
-    fn from_iter<T: Iterator<A>>(iterator: T) -> DList<A> {
+    fn from_iter_with_capacity<T: Iterator<A>>(iterator: T, _cap: uint) -> DList<A> {
         let mut ret = DList::new();
         ret.extend(iterator);
         ret
     }
 }
 
 impl<A> Extendable<A> for DList<A> {
-    fn extend<T: Iterator<A>>(&mut self, mut iterator: T) {
+    fn extend_with_capacity<T: Iterator<A>>(&mut self, mut iterator: T, _extra: uint) {
         for elt in iterator { self.push_back(elt); }
     }
 }

src/libcollections/hashmap.rs

@@ -1082,6 +1082,17 @@ impl<K: TotalEq + Hash<S>, V, S, H: Hasher<S>> HashMap<K, V, H> {
         }
     }
 
+    /// Reserve space for an additional `n` elements in the hash table.
+    pub fn reserve_additional(&mut self, extra: uint) {
+        let len = self.len();
+        if self.minimum_capacity - len < extra {
+            match len.checked_add(&extra) {
+                None => fail!("HashMap::reserve_additional: `uint` overflow"),
+                Some(new_cap) => self.reserve(new_cap)
+            }
+        }
+    }
+
     /// Resizes the internal vectors to a new capacity. It's your responsibility to:
     ///   1) Make sure the new capacity is enough for all the elements, accounting
     ///      for the load factor.
@@ -1454,16 +1465,16 @@ pub type Values<'a, K, V> =
     iter::Map<'static, (&'a K, &'a V), &'a V, Entries<'a, K, V>>;
 
 impl<K: TotalEq + Hash<S>, V, S, H: Hasher<S> + Default> FromIterator<(K, V)> for HashMap<K, V, H> {
-    fn from_iter<T: Iterator<(K, V)>>(iter: T) -> HashMap<K, V, H> {
-        let (lower, _) = iter.size_hint();
-        let mut map = HashMap::with_capacity_and_hasher(lower, Default::default());
+    fn from_iter_with_capacity<T: Iterator<(K, V)>>(iter: T, cap: uint) -> HashMap<K, V, H> {
+        let mut map = HashMap::with_capacity_and_hasher(cap, Default::default());
         map.extend(iter);
         map
     }
 }
 
 impl<K: TotalEq + Hash<S>, V, S, H: Hasher<S> + Default> Extendable<(K, V)> for HashMap<K, V, H> {
-    fn extend<T: Iterator<(K, V)>>(&mut self, mut iter: T) {
+    fn extend_with_capacity<T: Iterator<(K, V)>>(&mut self, mut iter: T, extra: uint) {
+        self.reserve_additional(extra);
         for (k, v) in iter {
             self.insert(k, v);
         }
@@ -1554,6 +1565,11 @@ impl<T: TotalEq + Hash<S>, S, H: Hasher<S>> HashSet<T, H> {
         self.map.reserve(n)
     }
 
+    /// Reserve space for an additional `n` elements in the hash table.
+    pub fn reserve_additional(&mut self, n: uint) {
+        self.map.reserve_additional(n)
+    }
+
     /// Returns true if the hash set contains a value equivalent to the
     /// given query value.
     pub fn contains_equiv<Q: Hash<S> + Equiv<T>>(&self, value: &Q) -> bool {
@@ -1617,16 +1633,16 @@ impl<T: TotalEq + Hash<S> + fmt::Show, S, H: Hasher<S>> fmt::Show for HashSet<T,
 }
 
 impl<T: TotalEq + Hash<S>, S, H: Hasher<S> + Default> FromIterator<T> for HashSet<T, H> {
-    fn from_iter<I: Iterator<T>>(iter: I) -> HashSet<T, H> {
-        let (lower, _) = iter.size_hint();
-        let mut set = HashSet::with_capacity_and_hasher(lower, Default::default());
+    fn from_iter_with_capacity<I: Iterator<T>>(iter: I, cap: uint) -> HashSet<T, H> {
+        let mut set = HashSet::with_capacity_and_hasher(cap, Default::default());
         set.extend(iter);
         set
     }
 }
 
 impl<T: TotalEq + Hash<S>, S, H: Hasher<S> + Default> Extendable<T> for HashSet<T, H> {
-    fn extend<I: Iterator<T>>(&mut self, mut iter: I) {
+    fn extend_with_capacity<I: Iterator<T>>(&mut self, mut iter: I, extra: uint) {
+        self.reserve_additional(extra);
         for k in iter {
             self.insert(k);
         }

src/libcollections/priority_queue.rs

@@ -60,6 +60,12 @@ impl<T:Ord> PriorityQueue<T> {
         self.data.reserve(n)
     }
 
+    /// Reserve capacity for an additional n elements in the PriorityQueue.
+    /// Do nothing if the capacity is already sufficient.
+    pub fn reserve_additional(&mut self, n: uint) {
+        self.data.reserve_additional(n)
+    }
+
     /// Pop the greatest item from the queue - fails if empty
     pub fn pop(&mut self) -> T {
         let mut item = self.data.pop().unwrap();
@@ -198,20 +204,16 @@ impl<'a, T> Iterator<&'a T> for Items<'a, T> {
 }
 
 impl<T: Ord> FromIterator<T> for PriorityQueue<T> {
-    fn from_iter<Iter: Iterator<T>>(iter: Iter) -> PriorityQueue<T> {
-        let mut q = PriorityQueue::new();
+    fn from_iter_with_capacity<Iter: Iterator<T>>(iter: Iter, cap: uint) -> PriorityQueue<T> {
+        let mut q = PriorityQueue::with_capacity(cap);
         q.extend(iter);
         q
     }
 }
 
 impl<T: Ord> Extendable<T> for PriorityQueue<T> {
-    fn extend<Iter: Iterator<T>>(&mut self, mut iter: Iter) {
-        let (lower, _) = iter.size_hint();
-
-        let len = self.capacity();
-        self.reserve(len + lower);
-
+    fn extend_with_capacity<Iter: Iterator<T>>(&mut self, mut iter: Iter, extra: uint) {
+        self.reserve_additional(extra);
         for elem in iter {
             self.push(elem);
         }

src/libcollections/ringbuf.rs

@@ -185,6 +185,20 @@ impl<T> RingBuf<T> {
         self.elts.reserve(n);
     }
 
+    /// Reserve capacity for an additional `n` elements in the given RingBuf,
+    /// over-allocating in case the caller needs to reserve additional
+    /// space.
+    ///
+    /// Do nothing if `self`'s capacity is already equal to or greater
+    /// than the requested capacity.
+    ///
+    /// # Arguments
+    ///
+    /// * n - The number of elements to reserve space for
+    pub fn reserve_additional(&mut self, n: uint) {
+        self.elts.reserve_additional(n);
+    }
+
     /// Front-to-back iterator.
     pub fn iter<'a>(&'a self) -> Items<'a, T> {
         Items{index: 0, rindex: self.nelts, lo: self.lo, elts: self.elts.as_slice()}
@@ -385,16 +399,16 @@ impl<A: Eq> Eq for RingBuf<A> {
 }
 
 impl<A> FromIterator<A> for RingBuf<A> {
-    fn from_iter<T: Iterator<A>>(iterator: T) -> RingBuf<A> {
-        let (lower, _) = iterator.size_hint();
-        let mut deq = RingBuf::with_capacity(lower);
+    fn from_iter_with_capacity<T: Iterator<A>>(iterator: T, cap: uint) -> RingBuf<A> {
+        let mut deq = RingBuf::with_capacity(cap);
         deq.extend(iterator);
         deq
     }
 }
 
 impl<A> Extendable<A> for RingBuf<A> {
-    fn extend<T: Iterator<A>>(&mut self, mut iterator: T) {
+    fn extend_with_capacity<T: Iterator<A>>(&mut self, mut iterator: T, extra: uint) {
+        self.reserve_additional(extra);
         for elt in iterator {
             self.push_back(elt);
         }

src/libcollections/treemap.rs

@@ -925,7 +925,7 @@ fn remove<K: TotalOrd, V>(node: &mut Option<Box<TreeNode<K, V>>>,
 }
 
 impl<K: TotalOrd, V> FromIterator<(K, V)> for TreeMap<K, V> {
-    fn from_iter<T: Iterator<(K, V)>>(iter: T) -> TreeMap<K, V> {
+    fn from_iter_with_capacity<T: Iterator<(K, V)>>(iter: T, _cap: uint) -> TreeMap<K, V> {
         let mut map = TreeMap::new();
         map.extend(iter);
         map
@@ -934,15 +934,15 @@ impl<K: TotalOrd, V> FromIterator<(K, V)> for TreeMap<K, V> {
 
 impl<K: TotalOrd, V> Extendable<(K, V)> for TreeMap<K, V> {
     #[inline]
-    fn extend<T: Iterator<(K, V)>>(&mut self, mut iter: T) {
+    fn extend_with_capacity<T: Iterator<(K, V)>>(&mut self, mut iter: T, _extra: uint) {
         for (k, v) in iter {
             self.insert(k, v);
         }
     }
 }
 
 impl<T: TotalOrd> FromIterator<T> for TreeSet<T> {
-    fn from_iter<Iter: Iterator<T>>(iter: Iter) -> TreeSet<T> {
+    fn from_iter_with_capacity<Iter: Iterator<T>>(iter: Iter, _cap: uint) -> TreeSet<T> {
         let mut set = TreeSet::new();
         set.extend(iter);
         set
@@ -951,7 +951,7 @@ impl<T: TotalOrd> FromIterator<T> for TreeSet<T> {
 
 impl<T: TotalOrd> Extendable<T> for TreeSet<T> {
     #[inline]
-    fn extend<Iter: Iterator<T>>(&mut self, mut iter: Iter) {
+    fn extend_with_capacity<Iter: Iterator<T>>(&mut self, mut iter: Iter, _extra: uint) {
         for elem in iter {
             self.insert(elem);
         }

src/libcollections/trie.rs

@@ -261,15 +261,15 @@ impl<T> TrieMap<T> {
 }
 
 impl<T> FromIterator<(uint, T)> for TrieMap<T> {
-    fn from_iter<Iter: Iterator<(uint, T)>>(iter: Iter) -> TrieMap<T> {
+    fn from_iter_with_capacity<Iter: Iterator<(uint, T)>>(iter: Iter, _cap: uint) -> TrieMap<T> {
         let mut map = TrieMap::new();
         map.extend(iter);
         map
     }
 }
 
 impl<T> Extendable<(uint, T)> for TrieMap<T> {
-    fn extend<Iter: Iterator<(uint, T)>>(&mut self, mut iter: Iter) {
+    fn extend_with_capacity<Iter: Iterator<(uint, T)>>(&mut self, mut iter: Iter, _extra: uint) {
         for (k, v) in iter {
             self.insert(k, v);
         }
@@ -360,15 +360,15 @@ impl TrieSet {
 }
 
 impl FromIterator<uint> for TrieSet {
-    fn from_iter<Iter: Iterator<uint>>(iter: Iter) -> TrieSet {
+    fn from_iter_with_capacity<Iter: Iterator<uint>>(iter: Iter, _cap: uint) -> TrieSet {
         let mut set = TrieSet::new();
         set.extend(iter);
         set
     }
 }
 
 impl Extendable<uint> for TrieSet {
-    fn extend<Iter: Iterator<uint>>(&mut self, mut iter: Iter) {
+    fn extend_with_capacity<Iter: Iterator<uint>>(&mut self, mut iter: Iter, _extra: uint) {
         for elem in iter {
             self.insert(elem);
         }

src/libcore/iter.rs

@@ -76,13 +76,30 @@ use mem;
 /// Conversion from an `Iterator`
 pub trait FromIterator<A> {
     /// Build a container with elements from an external iterator.
-    fn from_iter<T: Iterator<A>>(iterator: T) -> Self;
+    #[inline]
+    fn from_iter<T: Iterator<A>>(iterator: T) -> Self {
+        // Use the lower bound of the size as a hint by default.
+        let (lower, _) = iterator.size_hint();
+        FromIterator::from_iter_with_capacity(iterator, lower)
+    }
+
+    /// Build a container with elements from an external iterator with an
+    /// explicit capacity hint.
+    fn from_iter_with_capacity<T: Iterator<A>>(iterator: T, capacity: uint) -> Self;
 }
 
 /// A type growable from an `Iterator` implementation
 pub trait Extendable<A>: FromIterator<A> {
     /// Extend a container with the elements yielded by an iterator
-    fn extend<T: Iterator<A>>(&mut self, iterator: T);
+    fn extend<T: Iterator<A>>(&mut self, iterator: T) {
+        // Use the lower bound of the size as a hint by default.
+        let (lower, _) = iterator.size_hint();
+        self.extend_with_capacity(iterator, lower)
+    }
+
+    /// Extend a container with the elements yielded by an iterator with an
+    /// explicit capacity hint.
+    fn extend_with_capacity<T: Iterator<A>>(&mut self, iterator: T, extra: uint);
 }
 
 /// An interface for dealing with "external iterators". These types of iterators
@@ -356,6 +373,30 @@ pub trait Iterator<A> {
         FlatMap{iter: self, f: f, frontiter: None, backiter: None }
     }
 
+    /// Creates an iterator that allows a closure to consume more than one
+    /// element at a time from the iterator.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use std::iter::count;
+    ///
+    /// let mut it = count(0, 1).batch(|iter| {
+    ///     Some(iter.by_ref().take(3).collect::<Vec<int>>())
+    /// }).take(3);
+    /// assert_eq!(it.next().unwrap(), vec!(0, 1, 2));
+    /// assert_eq!(it.next().unwrap(), vec!(3, 4, 5));
+    /// assert_eq!(it.next().unwrap(), vec!(6, 7, 8));
+    /// assert!(it.next().is_none());
+    /// ```
+    #[inline]
+    fn batch<'r, B>(self, f: |&mut Self|: 'r -> Option<B>) -> Batch<'r, A, B, Self> {
+        Batch {
+            iter: self,
+            f: f,
+        }
+    }
+
     /// Creates an iterator that yields `None` forever after the underlying
     /// iterator yields `None`. Random-access iterator behavior is not
     /// affected, only single and double-ended iterator behavior.
@@ -463,6 +504,22 @@ pub trait Iterator<A> {
         FromIterator::from_iter(self.by_ref())
     }
 
+    /// Loops through the entire iterator, collecting all of the elements into
+    /// a container implementing `FromIterator` with an explicit capacity hint.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// let a = [1, 2, 3, 4, 5];
+    /// let b: Vec<int> = a.iter().map(|&x| x).collect_with_capacity(10);
+    /// assert!(a.as_slice() == b.as_slice());
+    /// assert!(a.capacity() == 10);
+    /// ```
+    #[inline]
+    fn collect_with_capacity<B: FromIterator<A>>(&mut self, capacity: uint) -> B {
+        FromIterator::from_iter_with_capacity(self.by_ref(), capacity)
+    }
+
     /// Loops through `n` iterations, returning the `n`th element of the
     /// iterator.
     ///
@@ -1742,6 +1799,25 @@ impl<'a,
     }
 }
 
+/// An iterator that batches elements of `iter` with `predicate`.
+pub struct Batch<'r, A, B, T> {
+    iter: T,
+    f: |&mut T|: 'r -> Option<B>,
+}
+
+impl<'r, A, B, T: Iterator<A>> Iterator<B> for Batch<'r, A, B, T> {
+    #[inline]
+    fn next(&mut self) -> Option<B> {
+        (self.f)(&mut self.iter)
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        let (_, upper) = self.iter.size_hint();
+        (0, upper) // can't know a lower bound, due to the predicate
+    }
+}
+
 /// An iterator that yields `None` forever after the underlying iterator
 /// yields `None` once.
 #[deriving(Clone)]
@@ -2340,8 +2416,8 @@ mod tests {
     use uint;
 
     impl<T> FromIterator<T> for Vec<T> {
-        fn from_iter<I: Iterator<T>>(mut iterator: I) -> Vec<T> {
-            let mut v = Vec::new();
+        fn from_iter_with_capacity<I: Iterator<T>>(mut iterator: I, cap: uint) -> Vec<T> {
+            let mut v = Vec::with_capacity(cap);
             for e in iterator {
                 v.push(e);
             }
@@ -2403,6 +2479,26 @@ mod tests {
         assert!(it.collect::<Vec<uint>>() == vec![0*0, 2*2, 4*4, 6*6, 8*8]);
     }
 
+    #[test]
+    fn test_iterator_batch() {
+        let mut it = count(0u, 1u).batch(|it| {
+            Some(it.by_ref().take(3).collect::<Vec<uint>>())
+        }).take(3);
+
+        assert!(it.next() == Some(vec!(0, 1, 2)));
+        assert!(it.next() == Some(vec!(3, 4, 5)));
+        assert!(it.next() == Some(vec!(6, 7, 8)));
+        assert!(it.next().is_none());
+
+        let mut it = count(0u, 1u).take(10).batch(|it| {
+            match it.next() {
+                Some(x) if x < 5 => Some(x),
+                _ => None
+            }
+        });
+        assert!(it.collect::<Vec<uint>>() == vec!(0, 1, 2, 3, 4));
+    }
+
     #[test]
     fn test_iterator_enumerate() {
         let xs = [0u, 1, 2, 3, 4, 5];