auto merge of #8604 : kballard/rust/iter-size-hint, r=graydon

Implement `size_hint()` on the new std::vec Iterators.

Add or update `size_hint()` on std::iterator Iterators where appropriate.

r? @thestinger

2 files changed, 100 insertions, 11 deletions

diff --git a/src/libstd/iterator.rs b/src/libstd/iterator.rs
index 36e62856464..c2d2e62a3c8 100644
--- a/src/libstd/iterator.rs
+++ b/src/libstd/iterator.rs
@@ -18,7 +18,7 @@ implementing the `Iterator` trait.
 */
 
 use cmp;
-use num::{Zero, One, Integer, Saturating};
+use num::{Zero, One, Integer, CheckedAdd, Saturating};
 use option::{Option, Some, None};
 use ops::{Add, Mul, Sub};
 use cmp::Ord;
@@ -838,7 +838,7 @@ impl<A, T: Iterator<A>, U: Iterator<A>> Iterator<A> for Chain<T, U> {
         let lower = a_lower.saturating_add(b_lower);
 
         let upper = match (a_upper, b_upper) {
-            (Some(x), Some(y)) => Some(x.saturating_add(y)),
+            (Some(x), Some(y)) => x.checked_add(&y),
             _ => None
         };
 
@@ -1115,6 +1115,21 @@ impl<A, T: Iterator<A>> Iterator<A> for Peekable<A, T> {
         if self.peeked.is_some() { self.peeked.take() }
         else { self.iter.next() }
     }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        let (lo, hi) = self.iter.size_hint();
+        if self.peeked.is_some() {
+            let lo = lo.saturating_add(1);
+            let hi = match hi {
+                Some(x) => x.checked_add(&1),
+                None => None
+            };
+            (lo, hi)
+        } else {
+            (lo, hi)
+        }
+    }
 }
 
 impl<'self, A, T: Iterator<A>> Peekable<A, T> {
@@ -1122,15 +1137,12 @@ impl<'self, A, T: Iterator<A>> Peekable<A, T> {
     /// or None if the iterator is exhausted.
     #[inline]
     pub fn peek(&'self mut self) -> Option<&'self A> {
+        if self.peeked.is_none() {
+            self.peeked = self.iter.next();
+        }
         match self.peeked {
             Some(ref value) => Some(value),
-            None => {
-                self.peeked = self.iter.next();
-                match self.peeked {
-                    Some(ref value) => Some(value),
-                    None => None,
-                }
-            },
+            None => None,
         }
     }
 }
@@ -1376,7 +1388,7 @@ impl<'self, A, T: Iterator<A>, B, U: Iterator<B>> Iterator<B> for
         let (blo, bhi) = self.backiter.map_default((0, Some(0)), |it| it.size_hint());
         let lo = flo.saturating_add(blo);
         match (self.iter.size_hint(), fhi, bhi) {
-            ((0, Some(0)), Some(a), Some(b)) => (lo, Some(a.saturating_add(b))),
+            ((0, Some(0)), Some(a), Some(b)) => (lo, a.checked_add(&b)),
             _ => (lo, None)
         }
     }
@@ -1482,6 +1494,12 @@ impl<'self, A, St> Iterator<A> for Unfoldr<'self, A, St> {
     fn next(&mut self) -> Option<A> {
         (self.f)(&mut self.state)
     }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        // no possible known bounds at this point
+        (0, None)
+    }
 }
 
 /// An infinite iterator starting at `start` and advancing by `step` with each
@@ -1525,6 +1543,9 @@ impl<A: Add<A, A> + Ord + Clone> Iterator<A> for Range<A> {
             None
         }
     }
+
+    // FIXME: #8606 Implement size_hint() on Range
+    // Blocked on #8605 Need numeric trait for converting to `Option<uint>`
 }
 
 impl<A: Sub<A, A> + Integer + Ord + Clone> DoubleEndedIterator<A> for Range<A> {
diff --git a/src/libstd/vec.rs b/src/libstd/vec.rs
index 0427eec2b05..01e7e053cf5 100644
--- a/src/libstd/vec.rs
+++ b/src/libstd/vec.rs
@@ -65,7 +65,7 @@ use cmp::{Eq, TotalOrd, Ordering, Less, Equal, Greater};
 use cmp;
 use iterator::*;
 use libc::c_void;
-use num::Zero;
+use num::{Integer, Zero, CheckedAdd, Saturating};
 use option::{None, Option, Some};
 use ptr::to_unsafe_ptr;
 use ptr;
@@ -209,6 +209,7 @@ pub struct SplitIterator<'self, T> {
 }
 
 impl<'self, T> Iterator<&'self [T]> for SplitIterator<'self, T> {
+    #[inline]
     fn next(&mut self) -> Option<&'self [T]> {
         if self.finished { return None; }
 
@@ -230,6 +231,21 @@ impl<'self, T> Iterator<&'self [T]> for SplitIterator<'self, T> {
             }
         }
     }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        if self.finished {
+            return (0, Some(0))
+        }
+        // if the predicate doesn't match anything, we yield one slice
+        // if it matches every element, we yield N+1 empty slices where
+        // N is either the number of elements or the number of splits.
+        match (self.v.len(), self.n) {
+            (0,_) => (1, Some(1)),
+            (_,0) => (1, Some(1)),
+            (l,n) => (1, cmp::min(l,n).checked_add(&1u))
+        }
+    }
 }
 
 /// An iterator over the slices of a vector separated by elements that
@@ -242,6 +258,7 @@ pub struct RSplitIterator<'self, T> {
 }
 
 impl<'self, T> Iterator<&'self [T]> for RSplitIterator<'self, T> {
+    #[inline]
     fn next(&mut self) -> Option<&'self [T]> {
         if self.finished { return None; }
 
@@ -263,6 +280,18 @@ impl<'self, T> Iterator<&'self [T]> for RSplitIterator<'self, T> {
             }
         }
     }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        if self.finished {
+            return (0, Some(0))
+        }
+        match (self.v.len(), self.n) {
+            (0,_) => (1, Some(1)),
+            (_,0) => (1, Some(1)),
+            (l,n) => (1, cmp::min(l,n).checked_add(&1u))
+        }
+    }
 }
 
 // Appending
@@ -453,6 +482,7 @@ pub struct WindowIter<'self, T> {
 }
 
 impl<'self, T> Iterator<&'self [T]> for WindowIter<'self, T> {
+    #[inline]
     fn next(&mut self) -> Option<&'self [T]> {
         if self.size > self.v.len() {
             None
@@ -462,6 +492,16 @@ impl<'self, T> Iterator<&'self [T]> for WindowIter<'self, T> {
             ret
         }
     }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        if self.size > self.v.len() {
+            (0, Some(0))
+        } else {
+            let x = self.v.len() - self.size;
+            (x.saturating_add(1), x.checked_add(&1u))
+        }
+    }
 }
 
 /// An iterator over a vector in (non-overlapping) chunks (`size`
@@ -476,6 +516,7 @@ pub struct ChunkIter<'self, T> {
 }
 
 impl<'self, T> Iterator<&'self [T]> for ChunkIter<'self, T> {
+    #[inline]
     fn next(&mut self) -> Option<&'self [T]> {
         if self.v.len() == 0 {
             None
@@ -487,9 +528,21 @@ impl<'self, T> Iterator<&'self [T]> for ChunkIter<'self, T> {
             Some(fst)
         }
     }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        if self.v.len() == 0 {
+            (0, Some(0))
+        } else {
+            let (n, rem) = self.v.len().div_rem(&self.size);
+            let n = if rem > 0 { n+1 } else { n };
+            (n, Some(n))
+        }
+    }
 }
 
 impl<'self, T> DoubleEndedIterator<&'self [T]> for ChunkIter<'self, T> {
+    #[inline]
     fn next_back(&mut self) -> Option<&'self [T]> {
         if self.v.len() == 0 {
             None
@@ -2236,6 +2289,7 @@ impl<'self, T> RandomAccessIterator<&'self T> for VecIterator<'self, T> {
         exact
     }
 
+    #[inline]
     fn idx(&self, index: uint) -> Option<&'self T> {
         unsafe {
             if index < self.indexable() {
@@ -2281,6 +2335,7 @@ pub struct MoveIterator<T> {
 }
 
 impl<T> Iterator<T> for MoveIterator<T> {
+    #[inline]
     fn next(&mut self) -> Option<T> {
         // this is peculiar, but is required for safety with respect
         // to dtors. It traverses the first half of the vec, and
@@ -2298,6 +2353,12 @@ impl<T> Iterator<T> for MoveIterator<T> {
 
         self.v.pop_opt()
     }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        let l = self.v.len();
+        (l, Some(l))
+    }
 }
 
 /// An iterator that moves out of a vector in reverse order.
@@ -2307,9 +2368,16 @@ pub struct MoveRevIterator<T> {
 }
 
 impl<T> Iterator<T> for MoveRevIterator<T> {
+    #[inline]
     fn next(&mut self) -> Option<T> {
         self.v.pop_opt()
     }
+
+    #[inline]
+    fn size_hint(&self) -> (uint, Option<uint>) {
+        let l = self.v.len();
+        (l, Some(l))
+    }
 }
 
 impl<A> FromIterator<A> for ~[A] {