Added interleave_shortest adaptor.

Author: DanielKeep

src/adaptors.rs

@@ -70,6 +70,84 @@ impl<I, J> Iterator for Interleave<I, J> where
     }
 }
 
+/// An iterator adaptor that alternates elements from the two iterators until
+/// one of them runs out.
+///
+/// This iterator is *fused*.
+///
+/// See [*.interleave_shortest()*](trait.Itertools.html#method.interleave_shortest)
+/// for more information.
+#[derive(Clone)]
+pub struct InterleaveShortest<I, J> where
+    I: Iterator,
+    J: Iterator<Item=I::Item>,
+{
+    it0: Fuse<I>,
+    it1: Fuse<J>,
+    phase: bool, // false ==> it0, true ==> it1
+}
+
+impl<I, J> InterleaveShortest<I, J> where
+    I: Iterator,
+    J: Iterator<Item=I::Item>,
+{
+    /// Create a new **InterleaveShortest** iterator.
+    pub fn new(a: I, b: J) -> InterleaveShortest<I, J> {
+        InterleaveShortest {
+            it0: a.fuse(),
+            it1: b.fuse(),
+            phase: false,
+        }
+    }
+}
+
+impl<I, J> Iterator for InterleaveShortest<I, J> where
+    I: Iterator,
+    J: Iterator<Item=I::Item>,
+{
+    type Item = I::Item;
+
+    #[inline]
+    fn next(&mut self) -> Option<I::Item> {
+        match self.phase {
+            false => match self.it0.next() {
+                None => None,
+                e => {
+                    self.phase = true;
+                    e
+                }
+            },
+            true => match self.it1.next() {
+                None => None,
+                e => {
+                    self.phase = false;
+                    e
+                }
+            },
+        }
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        fn bound(a: usize, b: usize) -> Option<usize> {
+            use std::cmp::min;
+            2usize.checked_mul(min(a, b))
+                .and_then(|lhs| lhs.checked_add(if a > b { 1 } else { 0 }))
+        }
+
+        let (l0, u0) = self.it0.size_hint();
+        let (l1, u1) = self.it1.size_hint();
+        let lb = bound(l0, l1).unwrap_or(usize::max_value());
+        let ub = match (u0, u1) {
+            (None, None) => None,
+            (Some(u0), None) => 2usize.checked_mul(u0),
+            (None, Some(u1)) => 2usize.checked_mul(u1).and_then(|l| l.checked_add(1)),
+            (Some(u0), Some(u1)) => bound(u0, u1)
+        };
+        (lb, ub)
+    }
+}
+
 /// **Deprecated:** Use *.map_fn()* instead.
 pub struct FnMap<B, I> where
     I: Iterator,

src/lib.rs

@@ -39,6 +39,7 @@ use std::fmt;
 
 pub use adaptors::{
     Interleave,
+    InterleaveShortest,
     Product,
     PutBack,
     PutBackN,
@@ -240,6 +241,25 @@ pub trait Itertools : Iterator {
         Interleave::new(self, other.into_iter())
     }
 
+    /// Alternate elements from two iterators until one of them runs out.
+    ///
+    /// Iterator element type is **Self::Item**.
+    ///
+    /// This iterator is *fused*.
+    ///
+    /// ```
+    /// use itertools::Itertools;
+    ///
+    /// let it = (0..5).interleave_shortest(vec![7, 8]);
+    /// itertools::assert_equal(it, vec![0, 7, 1, 8, 2]);
+    /// ```
+    fn interleave_shortest<J>(self, other: J) -> InterleaveShortest<Self, J::IntoIter> where
+        J: IntoIterator<Item=Self::Item>,
+        Self: Sized
+    {
+        InterleaveShortest::new(self, other.into_iter())
+    }
+
     /// An iterator adaptor to insert a particular value
     /// between each element of the adapted iterator.
     ///

tests/tests.rs

@@ -96,6 +96,31 @@ fn interleave() {
     it::assert_equal(it, rs.iter());
 }
 
+#[test]
+fn interleave_shortest() {
+    let v0: Vec<i32> = vec![0, 2, 4];
+    let v1: Vec<i32> = vec![1, 3, 5, 7];
+    let it = v0.into_iter().interleave_shortest(v1.into_iter());
+    assert_eq!(it.size_hint(), (6, Some(6)));
+    assert_eq!(it.collect_vec(), vec![0, 1, 2, 3, 4, 5]);
+
+    let v0: Vec<i32> = vec![0, 2, 4, 6, 8];
+    let v1: Vec<i32> = vec![1, 3, 5];
+    let it = v0.into_iter().interleave_shortest(v1.into_iter());
+    assert_eq!(it.size_hint(), (7, Some(7)));
+    assert_eq!(it.collect_vec(), vec![0, 1, 2, 3, 4, 5, 6]);
+
+    let i0 = ::std::iter::repeat(0);
+    let v1: Vec<_> = vec![1, 3, 5];
+    let it = i0.interleave_shortest(v1.into_iter());
+    assert_eq!(it.size_hint(), (7, Some(7)));
+
+    let v0: Vec<_> = vec![0, 2, 4];
+    let i1 = ::std::iter::repeat(1);
+    let it = v0.into_iter().interleave_shortest(i1);
+    assert_eq!(it.size_hint(), (6, Some(6)));
+}
+
 #[test]
 fn times() {
     assert!(it::times(0).count() == 0);