auto merge of #7570 : kballard/rust/iterator-size-hint, r=thestinger

Change the signature of Iterator.size_hint() to always have a lower bound.

Implement .size_hint() on all remaining iterators (if it differs from the default).

6 files changed, 161 insertions, 31 deletions

diff --git a/src/libextra/priority_queue.rs b/src/libextra/priority_queue.rs
index 3d1ca4a9818..1f7ba9f6530 100644
--- a/src/libextra/priority_queue.rs
+++ b/src/libextra/priority_queue.rs
@@ -186,6 +186,9 @@ pub struct PriorityQueueIterator <'self, T> {
 impl<'self, T> Iterator<&'self T> for PriorityQueueIterator<'self, T> {
     #[inline]
     fn next(&mut self) -> Option<(&'self T)> { self.iter.next() }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) { self.iter.size_hint() }
 }
 
 #[cfg(test)]
diff --git a/src/libextra/treemap.rs b/src/libextra/treemap.rs
index 4216b5a6d1a..17970f158dd 100644
--- a/src/libextra/treemap.rs
+++ b/src/libextra/treemap.rs
@@ -198,14 +198,15 @@ impl<K: TotalOrd, V> TreeMap<K, V> {
     /// Get a lazy iterator over the key-value pairs in the map.
     /// Requires that it be frozen (immutable).
     pub fn iter<'a>(&'a self) -> TreeMapIterator<'a, K, V> {
-        TreeMapIterator{stack: ~[], node: &self.root}
+        TreeMapIterator{stack: ~[], node: &self.root, remaining: self.length}
     }
 }
 
 /// Lazy forward iterator over a map
 pub struct TreeMapIterator<'self, K, V> {
     priv stack: ~[&'self ~TreeNode<K, V>],
-    priv node: &'self Option<~TreeNode<K, V>>
+    priv node: &'self Option<~TreeNode<K, V>>,
+    priv remaining: uint
 }
 
 impl<'self, K, V> Iterator<(&'self K, &'self V)> for TreeMapIterator<'self, K, V> {
@@ -222,12 +223,18 @@ impl<'self, K, V> Iterator<(&'self K, &'self V)> for TreeMapIterator<'self, K, V
               None => {
                 let res = self.stack.pop();
                 self.node = &res.right;
+                self.remaining -= 1;
                 return Some((&res.key, &res.value));
               }
             }
         }
         None
     }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        (self.remaining, Some(self.remaining))
+    }
 }
 
 impl<'self, T> Iterator<&'self T> for TreeSetIterator<'self, T> {
diff --git a/src/librustc/util/enum_set.rs b/src/librustc/util/enum_set.rs
index f9bd7a3508e..3ce645e012b 100644
--- a/src/librustc/util/enum_set.rs
+++ b/src/librustc/util/enum_set.rs
@@ -125,9 +125,9 @@ impl<E:CLike> Iterator<E> for EnumSetIterator<E> {
         Some(elem)
     }
 
-    fn size_hint(&self) -> (Option<uint>, Option<uint>) {
-        let exact = Some(self.bits.population_count());
-        (exact, exact)
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        let exact = self.bits.population_count();
+        (exact, Some(exact))
     }
 }
 
diff --git a/src/libstd/iterator.rs b/src/libstd/iterator.rs
index 77befbf19aa..b164bcbd28b 100644
--- a/src/libstd/iterator.rs
+++ b/src/libstd/iterator.rs
@@ -26,6 +26,7 @@ use option::{Option, Some, None};
 use ops::{Add, Mul};
 use cmp::Ord;
 use clone::Clone;
+use uint;
 
 /// Conversion from an `Iterator`
 pub trait FromIterator<A, T: Iterator<A>> {
@@ -43,7 +44,7 @@ pub trait Iterator<A> {
     /// Return a lower bound and upper bound on the remaining length of the iterator.
     ///
     /// The common use case for the estimate is pre-allocating space to store the results.
-    fn size_hint(&self) -> (Option<uint>, Option<uint>) { (None, None) }
+    fn size_hint(&self) -> (uint, Option<uint>) { (0, None) }
 }
 
 /// Iterator adaptors provided for every `Iterator` implementation. The adaptor objects are also
@@ -684,18 +685,18 @@ impl<A, T: Iterator<A>, U: Iterator<A>> Iterator<A> for ChainIterator<A, T, U> {
     }
 
     #[inline]
-    fn size_hint(&self) -> (Option<uint>, Option<uint>) {
+    fn size_hint(&self) -> (uint, Option<uint>) {
         let (a_lower, a_upper) = self.a.size_hint();
         let (b_lower, b_upper) = self.b.size_hint();
 
-        let lower = match (a_lower, b_lower) {
-            (Some(x), Some(y)) => Some(x + y),
-            (Some(x), None) => Some(x),
-            (None, Some(y)) => Some(y),
-            (None, None) => None
+        let lower = if uint::max_value - a_lower < b_lower {
+            uint::max_value
+        } else {
+            a_lower + b_lower
         };
 
         let upper = match (a_upper, b_upper) {
+            (Some(x), Some(y)) if uint::max_value - x < y => Some(uint::max_value),
             (Some(x), Some(y)) => Some(x + y),
             _ => None
         };
@@ -719,6 +720,23 @@ impl<A, B, T: Iterator<A>, U: Iterator<B>> Iterator<(A, B)> for ZipIterator<A, T
             _ => None
         }
     }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        let (a_lower, a_upper) = self.a.size_hint();
+        let (b_lower, b_upper) = self.b.size_hint();
+
+        let lower = cmp::min(a_lower, b_lower);
+
+        let upper = match (a_upper, b_upper) {
+            (Some(x), Some(y)) => Some(cmp::min(x,y)),
+            (Some(x), None) => Some(x),
+            (None, Some(y)) => Some(y),
+            (None, None) => None
+        };
+
+        (lower, upper)
+    }
 }
 
 /// An iterator which maps the values of `iter` with `f`
@@ -737,7 +755,7 @@ impl<'self, A, B, T: Iterator<A>> Iterator<B> for MapIterator<'self, A, B, T> {
     }
 
     #[inline]
-    fn size_hint(&self) -> (Option<uint>, Option<uint>) {
+    fn size_hint(&self) -> (uint, Option<uint>) {
         self.iter.size_hint()
     }
 }
@@ -762,9 +780,9 @@ impl<'self, A, T: Iterator<A>> Iterator<A> for FilterIterator<'self, A, T> {
     }
 
     #[inline]
-    fn size_hint(&self) -> (Option<uint>, Option<uint>) {
+    fn size_hint(&self) -> (uint, Option<uint>) {
         let (_, upper) = self.iter.size_hint();
-        (None, upper) // can't know a lower bound, due to the predicate
+        (0, upper) // can't know a lower bound, due to the predicate
     }
 }
 
@@ -787,9 +805,9 @@ impl<'self, A, B, T: Iterator<A>> Iterator<B> for FilterMapIterator<'self, A, B,
     }
 
     #[inline]
-    fn size_hint(&self) -> (Option<uint>, Option<uint>) {
+    fn size_hint(&self) -> (uint, Option<uint>) {
         let (_, upper) = self.iter.size_hint();
-        (None, upper) // can't know a lower bound, due to the predicate
+        (0, upper) // can't know a lower bound, due to the predicate
     }
 }
 
@@ -812,6 +830,11 @@ impl<A, T: Iterator<A>> Iterator<(uint, A)> for EnumerateIterator<A, T> {
             _ => None
         }
     }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        self.iter.size_hint()
+    }
 }
 
 /// An iterator which rejects elements while `predicate` is true
@@ -844,6 +867,12 @@ impl<'self, A, T: Iterator<A>> Iterator<A> for SkipWhileIterator<'self, A, T> {
             }
         }
     }
+
+    #[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 which only accepts elements while `predicate` is true
@@ -872,6 +901,12 @@ impl<'self, A, T: Iterator<A>> Iterator<A> for TakeWhileIterator<'self, A, T> {
             }
         }
     }
+
+    #[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 which skips over `n` elements of `iter`.
@@ -905,6 +940,21 @@ impl<A, T: Iterator<A>> Iterator<A> for SkipIterator<A, T> {
             next
         }
     }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        let (lower, upper) = self.iter.size_hint();
+
+        let lower = if lower >= self.n { lower - self.n } else { 0 };
+
+        let upper = match upper {
+            Some(x) if x >= self.n => Some(x - self.n),
+            Some(_) => Some(0),
+            None => None
+        };
+
+        (lower, upper)
+    }
 }
 
 /// An iterator which only iterates over the first `n` iterations of `iter`.
@@ -925,6 +975,20 @@ impl<A, T: Iterator<A>> Iterator<A> for TakeIterator<A, T> {
             None
         }
     }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        let (lower, upper) = self.iter.size_hint();
+
+        let lower = cmp::min(lower, self.n);
+
+        let upper = match upper {
+            Some(x) if x < self.n => Some(x),
+            _ => Some(self.n)
+        };
+
+        (lower, upper)
+    }
 }
 
 /// An iterator to maintain state while iterating another iterator
@@ -941,6 +1005,12 @@ impl<'self, A, B, T: Iterator<A>, St> Iterator<B> for ScanIterator<'self, A, B,
     fn next(&mut self) -> Option<B> {
         self.iter.next().chain(|a| (self.f)(&mut self.state, a))
     }
+
+    #[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 scan function
+    }
 }
 
 /// An iterator that maps each element to an iterator,
@@ -1022,6 +1092,11 @@ impl<A: Add<A, A> + Clone> Iterator<A> for Counter<A> {
         self.state = self.state.add(&self.step); // FIXME: #6050
         Some(result)
     }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        (uint::max_value, None) // Too bad we can't specify an infinite lower bound
+    }
 }
 
 #[cfg(test)]
@@ -1238,6 +1313,43 @@ mod tests {
     }
 
     #[test]
+    fn test_iterator_size_hint() {
+        let c = Counter::new(0, 1);
+        let v = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
+        let v2 = &[10, 11, 12];
+        let vi = v.iter();
+
+        assert_eq!(c.size_hint(), (uint::max_value, None));
+        assert_eq!(vi.size_hint(), (10, Some(10)));
+
+        assert_eq!(c.take_(5).size_hint(), (5, Some(5)));
+        assert_eq!(c.skip(5).size_hint().second(), None);
+        assert_eq!(c.take_while(|_| false).size_hint(), (0, None));
+        assert_eq!(c.skip_while(|_| false).size_hint(), (0, None));
+        assert_eq!(c.enumerate().size_hint(), (uint::max_value, None));
+        assert_eq!(c.chain_(vi.transform(|&i| i)).size_hint(), (uint::max_value, None));
+        assert_eq!(c.zip(vi).size_hint(), (10, Some(10)));
+        assert_eq!(c.scan(0, |_,_| Some(0)).size_hint(), (0, None));
+        assert_eq!(c.filter(|_| false).size_hint(), (0, None));
+        assert_eq!(c.transform(|_| 0).size_hint(), (uint::max_value, None));
+        assert_eq!(c.filter_map(|_| Some(0)).size_hint(), (0, None));
+
+        assert_eq!(vi.take_(5).size_hint(), (5, Some(5)));
+        assert_eq!(vi.take_(12).size_hint(), (10, Some(10)));
+        assert_eq!(vi.skip(3).size_hint(), (7, Some(7)));
+        assert_eq!(vi.skip(12).size_hint(), (0, Some(0)));
+        assert_eq!(vi.take_while(|_| false).size_hint(), (0, Some(10)));
+        assert_eq!(vi.skip_while(|_| false).size_hint(), (0, Some(10)));
+        assert_eq!(vi.enumerate().size_hint(), (10, Some(10)));
+        assert_eq!(vi.chain_(v2.iter()).size_hint(), (13, Some(13)));
+        assert_eq!(vi.zip(v2.iter()).size_hint(), (3, Some(3)));
+        assert_eq!(vi.scan(0, |_,_| Some(0)).size_hint(), (0, Some(10)));
+        assert_eq!(vi.filter(|_| false).size_hint(), (0, Some(10)));
+        assert_eq!(vi.transform(|i| i+1).size_hint(), (10, Some(10)));
+        assert_eq!(vi.filter_map(|_| Some(0)).size_hint(), (0, Some(10)));
+    }
+
+    #[test]
     fn test_collect() {
         let a = ~[1, 2, 3, 4, 5];
         let b: ~[int] = a.iter().transform(|&x| x).collect();
diff --git a/src/libstd/vec.rs b/src/libstd/vec.rs
index 1014ff48b1d..7244a9a7aac 100644
--- a/src/libstd/vec.rs
+++ b/src/libstd/vec.rs
@@ -2024,14 +2024,14 @@ macro_rules! iterator {
             }
 
             #[inline]
-            fn size_hint(&self) -> (Option<uint>, Option<uint>) {
+            fn size_hint(&self) -> (uint, Option<uint>) {
                 let diff = if $step > 0 {
                     (self.end as uint) - (self.ptr as uint)
                 } else {
                     (self.ptr as uint) - (self.end as uint)
                 };
-                let exact = Some(diff / size_of::<$elem>());
-                (exact, exact)
+                let exact = diff / size_of::<$elem>();
+                (exact, Some(exact))
             }
         }
     }
@@ -2132,7 +2132,7 @@ impl<A, T: Iterator<A>> FromIterator<A, T> for ~[A] {
 impl<A, T: Iterator<A>> FromIterator<A, T> for ~[A] {
     pub fn from_iterator(iterator: &mut T) -> ~[A] {
         let (lower, _) = iterator.size_hint();
-        let mut xs = with_capacity(lower.get_or_zero());
+        let mut xs = with_capacity(lower);
         for iterator.advance |x| {
             xs.push(x);
         }
@@ -2968,17 +2968,17 @@ mod tests {
         use iterator::*;
         let xs = [1, 2, 5, 10, 11];
         let mut it = xs.iter();
-        assert_eq!(it.size_hint(), (Some(5), Some(5)));
+        assert_eq!(it.size_hint(), (5, Some(5)));
         assert_eq!(it.next().unwrap(), &1);
-        assert_eq!(it.size_hint(), (Some(4), Some(4)));
+        assert_eq!(it.size_hint(), (4, Some(4)));
         assert_eq!(it.next().unwrap(), &2);
-        assert_eq!(it.size_hint(), (Some(3), Some(3)));
+        assert_eq!(it.size_hint(), (3, Some(3)));
         assert_eq!(it.next().unwrap(), &5);
-        assert_eq!(it.size_hint(), (Some(2), Some(2)));
+        assert_eq!(it.size_hint(), (2, Some(2)));
         assert_eq!(it.next().unwrap(), &10);
-        assert_eq!(it.size_hint(), (Some(1), Some(1)));
+        assert_eq!(it.size_hint(), (1, Some(1)));
         assert_eq!(it.next().unwrap(), &11);
-        assert_eq!(it.size_hint(), (Some(0), Some(0)));
+        assert_eq!(it.size_hint(), (0, Some(0)));
         assert!(it.next().is_none());
     }
 
@@ -2986,10 +2986,10 @@ mod tests {
     fn test_iter_size_hints() {
         use iterator::*;
         let mut xs = [1, 2, 5, 10, 11];
-        assert_eq!(xs.iter().size_hint(), (Some(5), Some(5)));
-        assert_eq!(xs.rev_iter().size_hint(), (Some(5), Some(5)));
-        assert_eq!(xs.mut_iter().size_hint(), (Some(5), Some(5)));
-        assert_eq!(xs.mut_rev_iter().size_hint(), (Some(5), Some(5)));
+        assert_eq!(xs.iter().size_hint(), (5, Some(5)));
+        assert_eq!(xs.rev_iter().size_hint(), (5, Some(5)));
+        assert_eq!(xs.mut_iter().size_hint(), (5, Some(5)));
+        assert_eq!(xs.mut_rev_iter().size_hint(), (5, Some(5)));
     }
 
     #[test]
diff --git a/src/libsyntax/opt_vec.rs b/src/libsyntax/opt_vec.rs
index bf8c5ae462b..8e2da3d6eb1 100644
--- a/src/libsyntax/opt_vec.rs
+++ b/src/libsyntax/opt_vec.rs
@@ -146,4 +146,12 @@ impl<'self, T> Iterator<&'self T> for OptVecIterator<'self, T> {
             None => None
         }
     }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        match self.iter {
+            Some(ref x) => x.size_hint(),
+            None => (0, Some(0))
+        }
+    }
 }