Add Iterator::zip_longest.

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<A>: Iterator<A> {
     ///
     /// ```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<A>: Iterator<A> {
         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<B, U: Iterator<B>>(self, other: U) -> ZipLongest<Self, U> {
+        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<A>> ExactSizeIterator<B> for Map<'a, A, B, T
 #[unstable = "trait is unstable"]
 impl<A, B, T, U> ExactSizeIterator<(A, B)> for Zip<T, U>
     where T: ExactSizeIterator<A>, U: ExactSizeIterator<B> {}
+#[unstable = "trait is unstable"]
+impl<A, B, T, U> ExactSizeIterator<EitherOrBoth<A, B>> for ZipLongest<T, U>
+    where T: ExactSizeIterator<A>, U: ExactSizeIterator<B> {}
 
 /// An double-ended iterator with the direction inverted
 #[deriving(Clone)]
@@ -1368,6 +1393,93 @@ RandomAccessIterator<(A, B)> for Zip<T, U> {
     }
 }
 
+/// An iterator which iterates two other iterators simultaneously
+#[deriving(Clone)]
+#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
+pub struct ZipLongest<T, U> {
+    a: T,
+    b: U
+}
+
+impl<A, B, T: Iterator<A>, U: Iterator<B>> Iterator<EitherOrBoth<A, B>> for ZipLongest<T, U> {
+    #[inline]
+    fn next(&mut self) -> Option<EitherOrBoth<A, B>> {
+        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<uint>) {
+        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<A, B, T: ExactSizeIterator<A>, U: ExactSizeIterator<B>> DoubleEndedIterator<EitherOrBoth<A, B>>
+for ZipLongest<T, U> {
+    #[inline]
+    fn next_back(&mut self) -> Option<EitherOrBoth<A, B>> {
+        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<A, B, T: RandomAccessIterator<A>, U: RandomAccessIterator<B>>
+RandomAccessIterator<EitherOrBoth<A, B>> for ZipLongest<T, U> {
+    #[inline]
+    fn indexable(&self) -> uint {
+        cmp::max(self.a.indexable(), self.b.indexable())
+    }
+
+    #[inline]
+    fn idx(&mut self, index: uint) -> Option<EitherOrBoth<A, B>> {
+        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<A, B> {
+    /// 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]

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];