diff --git a/src/libcore/iter.rs b/src/libcore/iter.rs index f9595f0663d57..f302898763502 100644 --- a/src/libcore/iter.rs +++ b/src/libcore/iter.rs @@ -64,6 +64,7 @@ use num::{ToPrimitive, Int}; use ops::{Add, Deref}; use option::{Option, Some, None}; use uint; +use self::EitherOrBoth::{Left, Right, Both}; #[deprecated = "renamed to Extend"] pub use self::Extend as Extendable; @@ -137,7 +138,7 @@ pub trait IteratorExt: Iterator { /// /// ```rust /// let a = [0i]; - /// let b = [1i]; + /// let b = [1i, 2i]; /// let mut it = a.iter().zip(b.iter()); /// let (x0, x1) = (0i, 1i); /// assert_eq!(it.next().unwrap(), (&x0, &x1)); @@ -149,6 +150,27 @@ pub trait IteratorExt: Iterator { Zip{a: self, b: other} } + /// Creates an iterator which iterates over both this and the specified + /// iterators simultaneously, yielding pairs of two optional elements. + /// When both iterators return None, all further invocations of next() will + /// return None. + /// + /// # Example + /// + /// ```rust + /// let a = [0i]; + /// let b = [1i, 2i]; + /// let mut it = a.iter().zip(b.iter()); + /// let (x0, x1, x2) = (0i, 1i, 2i); + /// assert_eq!(it.next().unwrap(), (Some(&x0), Some(&x1))); + /// assert_eq!(it.next().unwrap(), (None, Some(&x2))); + /// assert!(it.next().is_none()); + /// ``` + #[inline] + fn zip_longest>(self, other: U) -> ZipLongest { + ZipLongest{a: self, b: other} + } + /// Creates a new iterator which will apply the specified function to each /// element returned by the first, yielding the mapped element instead. /// @@ -780,6 +802,9 @@ impl<'a, A, B, T: ExactSizeIterator> ExactSizeIterator for Map<'a, A, B, T #[unstable = "trait is unstable"] impl ExactSizeIterator<(A, B)> for Zip where T: ExactSizeIterator, U: ExactSizeIterator {} +#[unstable = "trait is unstable"] +impl ExactSizeIterator> for ZipLongest + where T: ExactSizeIterator, U: ExactSizeIterator {} /// An double-ended iterator with the direction inverted #[deriving(Clone)] @@ -1368,6 +1393,93 @@ RandomAccessIterator<(A, B)> for Zip { } } +/// An iterator which iterates two other iterators simultaneously +#[deriving(Clone)] +#[must_use = "iterator adaptors are lazy and do nothing unless consumed"] +pub struct ZipLongest { + a: T, + b: U +} + +impl, U: Iterator> Iterator> for ZipLongest { + #[inline] + fn next(&mut self) -> Option> { + match (self.a.next(), self.b.next()) { + (None, None) => None, + (Some(a), None) => Some(Left(a)), + (None, Some(b)) => Some(Right(b)), + (Some(a), Some(b)) => Some(Both(a, b)), + } + } + + #[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::max(a_lower, b_lower); + + let upper = match (a_upper, b_upper) { + (Some(x), Some(y)) => Some(cmp::max(x,y)), + _ => None + }; + + (lower, upper) + } +} + +impl, U: ExactSizeIterator> DoubleEndedIterator> +for ZipLongest { + #[inline] + fn next_back(&mut self) -> Option> { + use cmp::{Equal, Greater, Less}; + match self.a.len().cmp(&self.b.len()) { + Equal => match (self.a.next_back(), self.b.next_back()) { + (None, None) => None, + (Some(a), Some(b)) => Some(Both(a, b)), + // These can only happen if .len() is inconsistent with .next_back() + (Some(a), None) => Some(Left(a)), + (None, Some(b)) => Some(Right(b)), + }, + Greater => self.a.next_back().map(Left), + Less => self.b.next_back().map(Right), + } + } +} + +impl, U: RandomAccessIterator> +RandomAccessIterator> for ZipLongest { + #[inline] + fn indexable(&self) -> uint { + cmp::max(self.a.indexable(), self.b.indexable()) + } + + #[inline] + fn idx(&mut self, index: uint) -> Option> { + match (self.a.idx(index), self.b.idx(index)) { + (None, None) => None, + (Some(a), None) => Some(Left(a)), + (None, Some(b)) => Some(Right(b)), + (Some(a), Some(b)) => Some(Both(a, b)), + } + } +} + +/// A value yielded by `ZipLongest`. +/// Contains one or two values, +/// depending on which of the input iterators are exhausted. +#[deriving(Clone, PartialEq, Eq, Show)] +pub enum EitherOrBoth { + /// Neither input iterator is exhausted yet, yielding two values. + Both(A, B), + /// The parameter iterator of `.zip_longest()` is exhausted, + /// only yielding a value from the `self` iterator. + Left(A), + /// The `self` iterator of `.zip_longest()` is exhausted, + /// only yielding a value from the parameter iterator. + Right(B), +} + /// An iterator which maps the values of `iter` with `f` #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] #[stable] diff --git a/src/libcoretest/iter.rs b/src/libcoretest/iter.rs index d046faa82d405..a1ce060978db2 100644 --- a/src/libcoretest/iter.rs +++ b/src/libcoretest/iter.rs @@ -340,6 +340,7 @@ fn test_iterator_size_hint() { assert_eq!(c.enumerate().size_hint(), (uint::MAX, None)); assert_eq!(c.chain(vi.map(|&i| i)).size_hint(), (uint::MAX, None)); assert_eq!(c.zip(vi).size_hint(), (10, Some(10))); + assert_eq!(c.zip_longest(vi).size_hint(), (uint::MAX, None)); assert_eq!(c.scan(0i, |_,_| Some(0i)).size_hint(), (0, None)); assert_eq!(c.filter(|_| false).size_hint(), (0, None)); assert_eq!(c.map(|_| 0i).size_hint(), (uint::MAX, None)); @@ -354,6 +355,7 @@ fn test_iterator_size_hint() { 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.zip_longest(v2.iter()).size_hint(), (10, Some(10))); assert_eq!(vi.scan(0i, |_,_| Some(0i)).size_hint(), (0, Some(10))); assert_eq!(vi.filter(|_| false).size_hint(), (0, Some(10))); assert_eq!(vi.map(|&i| i+1).size_hint(), (10, Some(10))); @@ -497,6 +499,23 @@ fn test_double_ended_zip() { assert_eq!(it.next(), None); } +#[test] +fn test_double_ended_zip_longest() { + use core::iter::EitherOrBoth::{Both, Left}; + let xs = [1i, 2, 3, 4, 5, 6]; + let ys = [1i, 2, 3, 7]; + let a = xs.iter().map(|&x| x); + let b = ys.iter().map(|&x| x); + let mut it = a.zip_longest(b); + assert_eq!(it.next(), Some(Both(1, 1))); + assert_eq!(it.next(), Some(Both(2, 2))); + assert_eq!(it.next_back(), Some(Left(6))); + assert_eq!(it.next_back(), Some(Left(5))); + assert_eq!(it.next_back(), Some(Both(4, 7))); + assert_eq!(it.next(), Some(Both(3, 3))); + assert_eq!(it.next(), None); +} + #[test] fn test_double_ended_filter() { let xs = [1i, 2, 3, 4, 5, 6]; @@ -641,6 +660,13 @@ fn test_random_access_zip() { check_randacc_iter(xs.iter().zip(ys.iter()), cmp::min(xs.len(), ys.len())); } +#[test] +fn test_random_access_zip_longest() { + let xs = [1i, 2, 3, 4, 5]; + let ys = [7i, 9, 11]; + check_randacc_iter(xs.iter().zip_longest(ys.iter()), cmp::max(xs.len(), ys.len())); +} + #[test] fn test_random_access_take() { let xs = [1i, 2, 3, 4, 5];