diff --git a/src/libstd/iterator.rs b/src/libstd/iterator.rs index 46d449e4dfb..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> { @@ -688,9 +689,14 @@ impl, U: Iterator> Iterator for ChainIterator { let (a_lower, a_upper) = self.a.size_hint(); let (b_lower, b_upper) = self.b.size_hint(); - let lower = a_lower + b_lower; + 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 }; @@ -714,6 +720,23 @@ impl, U: Iterator> Iterator<(A, B)> for ZipIterator None } } + + #[inline] + fn size_hint(&self) -> (uint, Option) { + 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` @@ -807,6 +830,11 @@ impl> Iterator<(uint, A)> for EnumerateIterator { _ => None } } + + #[inline] + fn size_hint(&self) -> (uint, Option) { + self.iter.size_hint() + } } /// An iterator which rejects elements while `predicate` is true @@ -839,6 +867,12 @@ impl<'self, A, T: Iterator> Iterator for SkipWhileIterator<'self, A, T> { } } } + + #[inline] + fn size_hint(&self) -> (uint, Option) { + 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 @@ -867,6 +901,12 @@ impl<'self, A, T: Iterator> Iterator for TakeWhileIterator<'self, A, T> { } } } + + #[inline] + fn size_hint(&self) -> (uint, Option) { + 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`. @@ -900,6 +940,21 @@ impl> Iterator for SkipIterator { next } } + + #[inline] + fn size_hint(&self) -> (uint, Option) { + 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`. @@ -920,6 +975,20 @@ impl> Iterator for TakeIterator { None } } + + #[inline] + fn size_hint(&self) -> (uint, Option) { + 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 @@ -936,6 +1005,12 @@ impl<'self, A, B, T: Iterator, St> Iterator for ScanIterator<'self, A, B, fn next(&mut self) -> Option { self.iter.next().chain(|a| (self.f)(&mut self.state, a)) } + + #[inline] + fn size_hint(&self) -> (uint, Option) { + 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, @@ -1017,6 +1092,11 @@ impl + Clone> Iterator for Counter { self.state = self.state.add(&self.step); // FIXME: #6050 Some(result) } + + #[inline] + fn size_hint(&self) -> (uint, Option) { + (uint::max_value, None) // Too bad we can't specify an infinite lower bound + } } #[cfg(test)] @@ -1232,6 +1312,43 @@ mod tests { assert_eq!(v.slice(0, 0).iter().transform(|&x| x).min(), None); } + #[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];