Skip to content

Commit

Permalink
zip_squash & zip_stretch
Browse files Browse the repository at this point in the history
Introduces new zip alternatives.
  • Loading branch information
JonathanWoollett-Light committed Dec 14, 2024
1 parent f80883b commit 1130e46
Show file tree
Hide file tree
Showing 4 changed files with 232 additions and 1 deletion.
2 changes: 2 additions & 0 deletions src/free.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
///
Expand Down
55 changes: 54 additions & 1 deletion src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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.
Expand Down
87 changes: 87 additions & 0 deletions src/zip_squash.rs
Original file line number Diff line number Diff line change
@@ -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,
{
}
89 changes: 89 additions & 0 deletions src/zip_stretch.rs
Original file line number Diff line number Diff line change
@@ -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,
{
}

0 comments on commit 1130e46

Please sign in to comment.