From 1130e46f57062f05567166d988338caefda29c54 Mon Sep 17 00:00:00 2001 From: Jonathan Woollett-Light Date: Sat, 14 Dec 2024 00:09:17 +0000 Subject: [PATCH] `zip_squash` & `zip_stretch` Introduces new zip alternatives. --- src/free.rs | 2 ++ src/lib.rs | 55 +++++++++++++++++++++++++++- src/zip_squash.rs | 87 ++++++++++++++++++++++++++++++++++++++++++++ src/zip_stretch.rs | 89 ++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 232 insertions(+), 1 deletion(-) create mode 100644 src/zip_squash.rs create mode 100644 src/zip_stretch.rs diff --git a/src/free.rs b/src/free.rs index 4c6820543..29bc197d8 100644 --- a/src/free.rs +++ b/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. /// diff --git a/src/lib.rs b/src/lib.rs index 834a48dea..2505fff76 100644 --- a/src/lib.rs +++ b/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(self, other: J) -> ZipSquash + where + J: 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(self, other: J) -> ZipStretch + where + J: IntoIterator, + ::IntoIter: ExactSizeIterator, + Self: ExactSizeIterator + Sized, + { + zip_stretch::zip_stretch(self, other) + } } impl 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. diff --git a/src/zip_squash.rs b/src/zip_squash.rs new file mode 100644 index 000000000..7ecd67dc6 --- /dev/null +++ b/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 { + 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: I, j: J) -> ZipSquash +where + I: IntoIterator, + J: IntoIterator, + ::IntoIter: ExactSizeIterator, + ::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 Iterator for ZipSquash +where + I: ExactSizeIterator, + J: ExactSizeIterator, +{ + type Item = (I::Item, J::Item); + + fn next(&mut self) -> Option { + 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) { + size_hint::min(self.a.size_hint(), self.b.size_hint()) + } +} + +impl ExactSizeIterator for ZipSquash +where + I: ExactSizeIterator, + J: ExactSizeIterator, +{ +} diff --git a/src/zip_stretch.rs b/src/zip_stretch.rs new file mode 100644 index 000000000..5dffcbbb1 --- /dev/null +++ b/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 { + 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: I, j: J) -> ZipStretch +where + I: IntoIterator, + J: IntoIterator, + ::IntoIter: ExactSizeIterator, + ::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 Iterator for ZipStretch +where + I: ExactSizeIterator, + J: ExactSizeIterator, +{ + type Item = EitherOrBoth; + + fn next(&mut self) -> Option { + 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) { + size_hint::min(self.a.size_hint(), self.b.size_hint()) + } +} + +impl ExactSizeIterator for ZipStretch +where + I: ExactSizeIterator, + J: ExactSizeIterator, +{ +}