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 15, 2024
1 parent f80883b commit 5e80e7c
Show file tree
Hide file tree
Showing 4 changed files with 254 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
57 changes: 56 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,61 @@ 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 cloning
/// elements in the shortest iterator.
///
/// Similar to [`Itertools::zip_longest`] except elements in the shortest
/// iterator are evenly spread.
///
/// ```
/// use itertools::Itertools;
/// 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![(1, 1),(1,2),(2,3)]);
/// ```
#[inline]
fn zip_stretch<J>(self, other: J) -> ZipStretch<Self, J::IntoIter>
where
J: IntoIterator,
<J as IntoIterator>::IntoIter: ExactSizeIterator,
<<J as IntoIterator>::IntoIter as IntoIterator>::Item: Clone,
Self: ExactSizeIterator + Sized,
<Self as Iterator>::Item: Clone,
{
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,
{
}
109 changes: 109 additions & 0 deletions src/zip_stretch.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
use super::size_hint;
use std::cmp::Ordering;
use std::fmt;

/// An iterator which iterates two other iterators simultaneously
/// always returning the first and last elements of both iterators by using
/// cloning to extend the length of the shortest iterator.
///
/// See [`.zip_stretch()`](crate::Itertools::zip_stretch) for more information.
#[derive(Clone)]
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct ZipStretch<I: ExactSizeIterator, J: ExactSizeIterator>
where
<I as Iterator>::Item: Clone,
<J as Iterator>::Item: Clone,
{
a: I,
b: J,
a_delta: f32,
b_delta: f32,
a_index: f32,
b_index: f32,
a_dupe: Option<<I as Iterator>::Item>,
b_dupe: Option<<J as Iterator>::Item>,
}

impl<I: ExactSizeIterator + fmt::Debug, J: ExactSizeIterator + fmt::Debug> fmt::Debug
for ZipStretch<I, J>
where
<I as Iterator>::Item: Clone,
<J as Iterator>::Item: Clone,
{
debug_fmt_fields!(ZipStretch, a, b, a_delta, b_delta, a_index, b_index);
}

/// Zips two iterators cloning elements 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,
<<I as IntoIterator>::IntoIter as IntoIterator>::Item: Clone,
<<J as IntoIterator>::IntoIter as IntoIterator>::Item: Clone,
{
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,
a_dupe: None,
b_dupe: None,
}
}

impl<I, J> Iterator for ZipStretch<I, J>
where
I: ExactSizeIterator,
J: ExactSizeIterator,
<I as Iterator>::Item: Clone,
<J as Iterator>::Item: Clone,
{
type Item = (I::Item, J::Item);

fn next(&mut self) -> Option<Self::Item> {
if self.a_index.fract() < self.a_delta {
self.a_dupe = self.a.next();
}
self.a_index += self.a_delta;

if self.b_index.fract() < self.b_delta {
self.b_dupe = self.b.next();
}
self.b_index += self.b_delta;

match (&self.a_dupe, &self.b_dupe) {
(Some(a), Some(b)) => Some((a.clone(), b.clone())),
(None, Some(_)) | (Some(_), None) => unreachable!(),
(None, None) => None,
}
}

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,
<I as Iterator>::Item: Clone,
<J as Iterator>::Item: Clone,
{
}

0 comments on commit 5e80e7c

Please sign in to comment.