Add Iterator::zip_longest.

Like `Iterator::zip`, but the returned iterator is as long as the longest
of the input iterators rather than shortest,
and yields `(Option<A>, Option<B>)` (which can not both be `None`)
rather than `(A, B)`.

An alternative could be to yield a new `EitherOrBoth` (?) enum
that is either `Both(A, B)`, `Left(A)`, or `Right(B)`.

Precedent: https://docs.python.org/library/itertools.html#itertools.izip_longest

Author: SimonSapin

src/libcore/iter.rs

@@ -137,7 +137,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 +149,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 +801,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<(Option<A>, Option<B>)> for ZipLongest<T, U>
+    where T: ExactSizeIterator<A>, U: ExactSizeIterator<B> {}
 
 /// An double-ended iterator with the direction inverted
 #[deriving(Clone)]
@@ -1368,6 +1392,71 @@ 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<(Option<A>, Option<B>)> for ZipLongest<T, U> {
+    #[inline]
+    fn next(&mut self) -> Option<(Option<A>, Option<B>)> {
+        match (self.a.next(), self.b.next()) {
+            (None, None) => None,
+            pair_of_options => Some(pair_of_options),
+        }
+    }
+
+    #[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: ExactSize<A>, U: ExactSize<B>> DoubleEndedIterator<(Option<A>, Option<B>)>
+for ZipLongest<T, U> {
+    #[inline]
+    fn next_back(&mut self) -> Option<(Option<A>, Option<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,
+                pair_of_options => Some(pair_of_options),
+            },
+            Greater => self.a.next_back().map(|x| (Some(x), None)),
+            Less => self.b.next_back().map(|y| (None, Some(y))),
+        }
+    }
+}
+
+impl<A, B, T: RandomAccessIterator<A>, U: RandomAccessIterator<B>>
+RandomAccessIterator<(Option<A>, Option<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<(Option<A>, Option<B>)> {
+        match (self.a.idx(index), self.b.idx(index)) {
+            (None, None) => None,
+            pair_of_options => Some(pair_of_options),
+        }
+    }
+}
+
 /// 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,22 @@ fn test_double_ended_zip() {
     assert_eq!(it.next(), None);
 }
 
+#[test]
+fn test_double_ended_zip_longest() {
+    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((Some(1), Some(1))));
+    assert_eq!(it.next(), Some((Some(2), Some(2))));
+    assert_eq!(it.next_back(), Some((Some(6), None)));
+    assert_eq!(it.next_back(), Some((Some(5), None)));
+    assert_eq!(it.next_back(), Some((Some(4), Some(7))));
+    assert_eq!(it.next(), Some((Some(3), Some(3))));
+    assert_eq!(it.next(), None);
+}
+
 #[test]
 fn test_double_ended_filter() {
     let xs = [1i, 2, 3, 4, 5, 6];
@@ -641,6 +659,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];