diff --git a/src/libcore/iterator.rs b/src/libcore/iterator.rs index 3811d28b030..685b4fcac7f 100644 --- a/src/libcore/iterator.rs +++ b/src/libcore/iterator.rs @@ -34,6 +34,7 @@ pub trait IteratorUtil { // FIXME: #5898: should be called map fn transform<'r, B>(self, f: &'r fn(A) -> B) -> MapIterator<'r, A, B, Self>; fn filter<'r>(self, predicate: &'r fn(&A) -> bool) -> FilterIterator<'r, A, Self>; + fn filter_map<'r, B>(self, f: &'r fn(A) -> Option) -> FilterMapIterator<'r, A, B, Self>; fn enumerate(self) -> EnumerateIterator; fn skip_while<'r>(self, predicate: &'r fn(&A) -> bool) -> SkipWhileIterator<'r, A, Self>; fn take_while<'r>(self, predicate: &'r fn(&A) -> bool) -> TakeWhileIterator<'r, A, Self>; @@ -74,6 +75,11 @@ impl> IteratorUtil for T { FilterIterator{iter: self, predicate: predicate} } + #[inline(always)] + fn filter_map<'r, B>(self, f: &'r fn(A) -> Option) -> FilterMapIterator<'r, A, B, T> { + FilterMapIterator { iter: self, f: f } + } + #[inline(always)] fn enumerate(self) -> EnumerateIterator { EnumerateIterator{iter: self, count: 0} @@ -213,6 +219,28 @@ impl<'self, A, T: Iterator> Iterator for FilterIterator<'self, A, T> { } } +pub struct FilterMapIterator<'self, A, B, T> { + priv iter: T, + priv f: &'self fn(A) -> Option +} + +impl<'self, A, B, T: Iterator> Iterator for FilterMapIterator<'self, A, B, T> { + #[inline] + fn next(&mut self) -> Option { + loop { + match self.iter.next() { + None => { return None; } + Some(a) => { + match (self.f)(a) { + Some(b) => { return Some(b); } + None => { loop; } + } + } + } + } + } +} + pub struct EnumerateIterator { priv iter: T, priv count: uint @@ -432,6 +460,13 @@ mod tests { assert_eq!(i, expected.len()); } + #[test] + fn test_filter_map() { + let it = Counter::new(0u, 1u).take(10) + .filter_map(|x: uint| if x.is_even() { Some(x*x) } else { None }); + assert_eq!(it.to_vec(), ~[0*0, 2*2, 4*4, 6*6, 8*8]); + } + #[test] fn test_iterator_enumerate() { let xs = [0u, 1, 2, 3, 4, 5];