zip_squash & zip_stretch

Introduces new zip alternatives.

Author: JonathanWoollett-Light

src/free.rs

@@ -28,6 +28,8 @@ pub use crate::put_back_n_impl::put_back_n;
 #[cfg(feature = "use_alloc")]
 pub use crate::rciter_impl::rciter;
 pub use crate::zip_eq_impl::zip_eq;
+pub use crate::zip_squash::zip_squash;
+pub use crate::zip_stretch::zip_stretch;
 
 /// Iterate `iterable` with a particular value inserted between each element.
 ///

src/lib.rs

@@ -144,6 +144,8 @@ pub mod structs {
     pub use crate::with_position::WithPosition;
     pub use crate::zip_eq_impl::ZipEq;
     pub use crate::zip_longest::ZipLongest;
+    pub use crate::zip_squash::ZipSquash;
+    pub use crate::zip_stretch::ZipStretch;
     pub use crate::ziptuple::Zip;
 }
 
@@ -235,6 +237,8 @@ mod unziptuple;
 mod with_position;
 mod zip_eq_impl;
 mod zip_longest;
+mod zip_squash;
+mod zip_stretch;
 mod ziptuple;
 
 #[macro_export]
@@ -4537,10 +4541,59 @@ pub trait Itertools: Iterator {
             _ => Err(sh),
         }
     }
+
+    /// Create an iterator which iterates over both this and the specified
+    /// iterator simultaneously, yielding pairs of elements.
+    ///
+    /// Similar to [`Iterator::zip`] except elements are evenly sampled from
+    /// the longest iterator.
+    ///
+    /// ```
+    /// use itertools::Itertools;
+    /// let a = vec![1, 2];
+    /// let b = vec![1, 2, 3];
+    ///
+    /// let it = a.into_iter().zip_squash(b.into_iter());
+    /// itertools::assert_equal(it, vec![(1, 1),(2,3)]);
+    /// ```
+    #[inline]
+    fn zip_squash<J>(self, other: J) -> ZipSquash<Self, J::IntoIter>
+    where
+        J: IntoIterator,
+        <J as IntoIterator>::IntoIter: ExactSizeIterator,
+        Self: ExactSizeIterator + Sized,
+    {
+        zip_squash::zip_squash(self, other)
+    }
+    /// Create an iterator which iterates over both this and the specified
+    /// iterator simultaneously, yielding pairs of elements.
+    ///
+    /// Always yielding the first and last elements of both iterators by using [`EitherOrBoth`].
+    ///
+    /// Similar to [`Itertools::zip_longest`] except elements in the shortest iterator are evenly
+    /// spread.
+    ///
+    /// ```
+    /// use itertools::Itertools;
+    /// use itertools::EitherOrBoth;
+    /// let a = vec![1, 2];
+    /// let b = vec![1, 2, 3];
+    ///
+    /// let it = a.into_iter().zip_stretch(b.into_iter());
+    /// itertools::assert_equal(it, vec![EitherOrBoth::Both(1, 1),EitherOrBoth::Right(2), EitherOrBoth::Both(2,3)]);
+    /// ```
+    #[inline]
+    fn zip_stretch<J>(self, other: J) -> ZipStretch<Self, J::IntoIter>
+    where
+        J: IntoIterator,
+        <J as IntoIterator>::IntoIter: ExactSizeIterator,
+        Self: ExactSizeIterator + Sized,
+    {
+        zip_stretch::zip_stretch(self, other)
+    }
 }
 
 impl<T> Itertools for T where T: Iterator + ?Sized {}
-
 /// Return `true` if both iterables produce equal sequences
 /// (elements pairwise equal and sequences of the same length),
 /// `false` otherwise.

src/zip_squash.rs

@@ -0,0 +1,87 @@
+use super::size_hint;
+use std::cmp::Ordering;
+
+/// An iterator which iterates two other iterators simultaneously
+/// always returning elements are evenly sampled from the longest iterator.
+///
+/// See [`.zip_squash()`](crate::Itertools::zip_squash) for more information.
+#[derive(Clone, Debug)]
+#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
+pub struct ZipSquash<I: ExactSizeIterator, J: ExactSizeIterator> {
+    a: I,
+    b: J,
+    a_delta: f32,
+    b_delta: f32,
+    a_index: f32,
+    b_index: f32,
+}
+
+/// Zips two iterators skipping elements of the longest iterator to ensure it fully consumes both
+/// iterators.
+///
+/// [`IntoIterator`] enabled version of [`Itertools::zip_squash`](crate::Itertools::zip_squash).
+pub fn zip_squash<I, J>(i: I, j: J) -> ZipSquash<I::IntoIter, J::IntoIter>
+where
+    I: IntoIterator,
+    J: IntoIterator,
+    <I as IntoIterator>::IntoIter: ExactSizeIterator,
+    <J as IntoIterator>::IntoIter: ExactSizeIterator,
+{
+    use std::iter::ExactSizeIterator;
+    let (a, b) = (i.into_iter(), j.into_iter());
+    let (a_delta, b_delta) = match a.len().cmp(&b.len()) {
+        Ordering::Equal => (1f32, 1f32),
+        Ordering::Less => (1f32, b.len() as f32 / a.len() as f32),
+        Ordering::Greater => (a.len() as f32 / b.len() as f32, 1f32),
+    };
+    debug_assert!(a_delta >= 1f32);
+    debug_assert!(b_delta >= 1f32);
+    ZipSquash {
+        a,
+        b,
+        a_delta,
+        b_delta,
+        a_index: 0f32,
+        b_index: 0f32,
+    }
+}
+
+impl<I, J> Iterator for ZipSquash<I, J>
+where
+    I: ExactSizeIterator,
+    J: ExactSizeIterator,
+{
+    type Item = (I::Item, J::Item);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let (a, b) = (self.a.next(), self.b.next());
+        let a_diff = (self.a_delta / (1f32 - self.a_index.fract())).ceil() as usize;
+        self.a_index += a_diff as f32 * self.a_delta;
+        if let Some(skip) = a_diff.checked_sub(2) {
+            self.a.nth(skip);
+        }
+
+        let b_diff = (self.b_delta / (1f32 - self.b_index.fract())).ceil() as usize;
+        self.b_index += b_diff as f32 * self.b_delta;
+        if let Some(skip) = b_diff.checked_sub(2) {
+            self.b.nth(skip);
+        }
+
+        match (a, b) {
+            (None, None) => None,
+            (Some(a), Some(b)) => Some((a, b)),
+            (None, Some(_)) | (Some(_), None) => unreachable!(),
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        size_hint::min(self.a.size_hint(), self.b.size_hint())
+    }
+}
+
+impl<I, J> ExactSizeIterator for ZipSquash<I, J>
+where
+    I: ExactSizeIterator,
+    J: ExactSizeIterator,
+{
+}

src/zip_stretch.rs

@@ -0,0 +1,89 @@
+use super::size_hint;
+use crate::either_or_both::EitherOrBoth;
+use std::cmp::Ordering;
+
+/// An iterator which iterates two other iterators simultaneously
+/// always returning the first and last elements of both iterators by using
+/// [`EitherOrBoth`] to extend the length of the shortest iterator.
+///
+/// See [`.zip_stretch()`](crate::Itertools::zip_stretch) for more information.
+#[derive(Clone, Debug)]
+#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
+pub struct ZipStretch<I: ExactSizeIterator, J: ExactSizeIterator> {
+    a: I,
+    b: J,
+    a_delta: f32,
+    b_delta: f32,
+    a_index: f32,
+    b_index: f32,
+}
+
+/// Zips two iterators using [`EitherOrBoth`] to extend the length of the shortest iterator to
+/// ensure it fully consumes both iterators.
+///
+/// [`IntoIterator`] enabled version of [`Itertools::zip_stretch`](crate::Itertools::zip_stretch).
+pub fn zip_stretch<I, J>(i: I, j: J) -> ZipStretch<I::IntoIter, J::IntoIter>
+where
+    I: IntoIterator,
+    J: IntoIterator,
+    <I as IntoIterator>::IntoIter: ExactSizeIterator,
+    <J as IntoIterator>::IntoIter: ExactSizeIterator,
+{
+    use std::iter::ExactSizeIterator;
+    let (a, b) = (i.into_iter(), j.into_iter());
+    let (a_delta, b_delta) = match a.len().cmp(&b.len()) {
+        Ordering::Equal => (1f32, 1f32),
+        Ordering::Less => (a.len() as f32 / b.len() as f32, 1f32),
+        Ordering::Greater => (1f32, b.len() as f32 / a.len() as f32),
+    };
+    debug_assert!(a_delta <= 1f32);
+    debug_assert!(b_delta <= 1f32);
+    ZipStretch {
+        a,
+        b,
+        a_delta,
+        b_delta,
+        a_index: 0f32,
+        b_index: 0f32,
+    }
+}
+
+impl<I, J> Iterator for ZipStretch<I, J>
+where
+    I: ExactSizeIterator,
+    J: ExactSizeIterator,
+{
+    type Item = EitherOrBoth<I::Item, J::Item>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let mut a_return = None;
+        if self.a_index.fract() < self.a_delta {
+            a_return = self.a.next();
+        }
+        self.a_index += self.a_delta;
+
+        let mut b_return = None;
+        if self.b_index.fract() < self.b_delta {
+            b_return = self.b.next();
+        }
+        self.b_index += self.b_delta;
+
+        match (a_return, b_return) {
+            (None, None) => None,
+            (Some(a), Some(b)) => Some(EitherOrBoth::Both(a, b)),
+            (None, Some(b)) => Some(EitherOrBoth::Right(b)),
+            (Some(a), None) => Some(EitherOrBoth::Left(a)),
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        size_hint::min(self.a.size_hint(), self.b.size_hint())
+    }
+}
+
+impl<I, J> ExactSizeIterator for ZipStretch<I, J>
+where
+    I: ExactSizeIterator,
+    J: ExactSizeIterator,
+{
+}