Skip to content

Commit b065292

Browse files
JonathanWoollett-LightJonathanWoollett-Light
committed
zip_squash & zip_stretch
Introduces new zip alternatives.
1 parent f80883b commit b065292

File tree

5 files changed

+292
-1
lines changed

5 files changed

+292
-1
lines changed

‎src/free.rs

+2
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ pub use crate::put_back_n_impl::put_back_n;
2828
#[cfg(feature = "use_alloc")]
2929
pub use crate::rciter_impl::rciter;
3030
pub use crate::zip_eq_impl::zip_eq;
31+
pub use crate::zip_squash::zip_squash;
32+
pub use crate::zip_stretch::zip_stretch;
3133

3234
/// Iterate `iterable` with a particular value inserted between each element.
3335
///

‎src/lib.rs

+56-1
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,8 @@ pub mod structs {
144144
pub use crate::with_position::WithPosition;
145145
pub use crate::zip_eq_impl::ZipEq;
146146
pub use crate::zip_longest::ZipLongest;
147+
pub use crate::zip_squash::ZipSquash;
148+
pub use crate::zip_stretch::ZipStretch;
147149
pub use crate::ziptuple::Zip;
148150
}
149151

@@ -235,6 +237,8 @@ mod unziptuple;
235237
mod with_position;
236238
mod zip_eq_impl;
237239
mod zip_longest;
240+
mod zip_squash;
241+
mod zip_stretch;
238242
mod ziptuple;
239243

240244
#[macro_export]
@@ -4537,10 +4541,61 @@ pub trait Itertools: Iterator {
45374541
_ => Err(sh),
45384542
}
45394543
}
4544+
4545+
/// Create an iterator which iterates over both this and the specified
4546+
/// iterator simultaneously, yielding pairs of elements.
4547+
///
4548+
/// Similar to [`Iterator::zip`] except elements are evenly sampled from
4549+
/// the longest iterator.
4550+
///
4551+
/// ```
4552+
/// use itertools::Itertools;
4553+
/// let a = vec![1, 2];
4554+
/// let b = vec![1, 2, 3];
4555+
///
4556+
/// let it = a.into_iter().zip_squash(b.into_iter());
4557+
/// itertools::assert_equal(it, vec![(1, 1), (2, 2)]);
4558+
/// ```
4559+
#[inline]
4560+
fn zip_squash<J>(self, other: J) -> ZipSquash<Self, J::IntoIter>
4561+
where
4562+
J: IntoIterator,
4563+
<J as IntoIterator>::IntoIter: ExactSizeIterator,
4564+
Self: ExactSizeIterator + Sized,
4565+
{
4566+
zip_squash::zip_squash(self, other)
4567+
}
4568+
/// Create an iterator which iterates over both this and the specified
4569+
/// iterator simultaneously, yielding pairs of elements.
4570+
///
4571+
/// Always yielding the first and last elements of both iterators by cloning
4572+
/// elements in the shortest iterator.
4573+
///
4574+
/// Similar to [`Itertools::zip_longest`] except elements in the shortest
4575+
/// iterator are evenly spread.
4576+
///
4577+
/// ```
4578+
/// use itertools::Itertools;
4579+
/// let a = vec![1, 2];
4580+
/// let b = vec![1, 2, 3];
4581+
///
4582+
/// let it = a.into_iter().zip_stretch(b.into_iter());
4583+
/// itertools::assert_equal(it, vec![(1, 1), (1, 2), (2, 3)]);
4584+
/// ```
4585+
#[inline]
4586+
fn zip_stretch<J>(self, other: J) -> ZipStretch<Self, J::IntoIter>
4587+
where
4588+
J: IntoIterator,
4589+
<J as IntoIterator>::IntoIter: ExactSizeIterator,
4590+
<<J as IntoIterator>::IntoIter as IntoIterator>::Item: Clone,
4591+
Self: ExactSizeIterator + Sized,
4592+
<Self as Iterator>::Item: Clone,
4593+
{
4594+
zip_stretch::zip_stretch(self, other)
4595+
}
45404596
}
45414597

45424598
impl<T> Itertools for T where T: Iterator + ?Sized {}
4543-
45444599
/// Return `true` if both iterables produce equal sequences
45454600
/// (elements pairwise equal and sequences of the same length),
45464601
/// `false` otherwise.

‎src/zip_squash.rs

+83
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
use super::size_hint;
2+
use std::cmp::Ordering;
3+
4+
/// An iterator which iterates two other iterators simultaneously
5+
/// always returning elements are evenly sampled from the longest iterator.
6+
///
7+
/// See [`.zip_squash()`](crate::Itertools::zip_squash) for more information.
8+
#[derive(Clone, Debug)]
9+
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
10+
pub struct ZipSquash<I: ExactSizeIterator, J: ExactSizeIterator> {
11+
a: I,
12+
b: J,
13+
a_delta: f32,
14+
b_delta: f32,
15+
a_index: f32,
16+
b_index: f32,
17+
}
18+
19+
/// Zips two iterators skipping elements of the longest iterator to ensure it fully consumes both
20+
/// iterators.
21+
///
22+
/// [`IntoIterator`] enabled version of [`Itertools::zip_squash`](crate::Itertools::zip_squash).
23+
pub fn zip_squash<I, J>(i: I, j: J) -> ZipSquash<I::IntoIter, J::IntoIter>
24+
where
25+
I: IntoIterator,
26+
J: IntoIterator,
27+
<I as IntoIterator>::IntoIter: ExactSizeIterator,
28+
<J as IntoIterator>::IntoIter: ExactSizeIterator,
29+
{
30+
use std::iter::ExactSizeIterator;
31+
let (a, b) = (i.into_iter(), j.into_iter());
32+
let (a_delta, b_delta) = match a.len().cmp(&b.len()) {
33+
Ordering::Equal => (1f32, 1f32),
34+
Ordering::Less => (1f32, b.len() as f32 / a.len() as f32),
35+
Ordering::Greater => (a.len() as f32 / b.len() as f32, 1f32),
36+
};
37+
debug_assert!(a_delta >= 1f32);
38+
debug_assert!(b_delta >= 1f32);
39+
ZipSquash {
40+
a,
41+
b,
42+
a_delta,
43+
b_delta,
44+
a_index: 0f32,
45+
b_index: 0f32,
46+
}
47+
}
48+
49+
impl<I, J> Iterator for ZipSquash<I, J>
50+
where
51+
I: ExactSizeIterator,
52+
J: ExactSizeIterator,
53+
{
54+
type Item = (I::Item, J::Item);
55+
56+
fn next(&mut self) -> Option<Self::Item> {
57+
let (a, b) = (self.a.next(), self.b.next());
58+
let a_new = self.a_index + self.a_delta;
59+
if let Some(skip) = ((a_new.floor() - self.a_index.floor()) as usize).checked_sub(2) {
60+
self.a.nth(skip);
61+
}
62+
self.a_index = a_new;
63+
64+
let b_new = self.b_index + self.b_delta;
65+
if let Some(skip) = ((b_new.floor() - self.b_index.floor()) as usize).checked_sub(2) {
66+
self.b.nth(skip);
67+
}
68+
self.b_index = b_new;
69+
70+
a.zip(b)
71+
}
72+
73+
fn size_hint(&self) -> (usize, Option<usize>) {
74+
size_hint::min(self.a.size_hint(), self.b.size_hint())
75+
}
76+
}
77+
78+
impl<I, J> ExactSizeIterator for ZipSquash<I, J>
79+
where
80+
I: ExactSizeIterator,
81+
J: ExactSizeIterator,
82+
{
83+
}

‎src/zip_stretch.rs

+105
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
use super::size_hint;
2+
use std::cmp::Ordering;
3+
use std::fmt;
4+
5+
/// An iterator which iterates two other iterators simultaneously
6+
/// always returning the first and last elements of both iterators by using
7+
/// cloning to extend the length of the shortest iterator.
8+
///
9+
/// See [`.zip_stretch()`](crate::Itertools::zip_stretch) for more information.
10+
#[derive(Clone)]
11+
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
12+
pub struct ZipStretch<I: ExactSizeIterator, J: ExactSizeIterator>
13+
where
14+
<I as Iterator>::Item: Clone,
15+
<J as Iterator>::Item: Clone,
16+
{
17+
a: I,
18+
b: J,
19+
a_delta: f32,
20+
b_delta: f32,
21+
a_index: f32,
22+
b_index: f32,
23+
a_dupe: Option<<I as Iterator>::Item>,
24+
b_dupe: Option<<J as Iterator>::Item>,
25+
}
26+
27+
impl<I: ExactSizeIterator + fmt::Debug, J: ExactSizeIterator + fmt::Debug> fmt::Debug
28+
for ZipStretch<I, J>
29+
where
30+
<I as Iterator>::Item: Clone,
31+
<J as Iterator>::Item: Clone,
32+
{
33+
debug_fmt_fields!(ZipStretch, a, b, a_delta, b_delta, a_index, b_index);
34+
}
35+
36+
/// Zips two iterators cloning elements to extend the length of the shortest iterator to
37+
/// ensure it fully consumes both iterators.
38+
///
39+
/// [`IntoIterator`] enabled version of [`Itertools::zip_stretch`](crate::Itertools::zip_stretch).
40+
pub fn zip_stretch<I, J>(i: I, j: J) -> ZipStretch<I::IntoIter, J::IntoIter>
41+
where
42+
I: IntoIterator,
43+
J: IntoIterator,
44+
<I as IntoIterator>::IntoIter: ExactSizeIterator,
45+
<J as IntoIterator>::IntoIter: ExactSizeIterator,
46+
<<I as IntoIterator>::IntoIter as IntoIterator>::Item: Clone,
47+
<<J as IntoIterator>::IntoIter as IntoIterator>::Item: Clone,
48+
{
49+
use std::iter::ExactSizeIterator;
50+
let (a, b) = (i.into_iter(), j.into_iter());
51+
let (a_delta, b_delta) = match a.len().cmp(&b.len()) {
52+
Ordering::Equal => (1f32, 1f32),
53+
Ordering::Less => (a.len() as f32 / b.len() as f32, 1f32),
54+
Ordering::Greater => (1f32, b.len() as f32 / a.len() as f32),
55+
};
56+
debug_assert!(a_delta <= 1f32);
57+
debug_assert!(b_delta <= 1f32);
58+
ZipStretch {
59+
a,
60+
b,
61+
a_delta,
62+
b_delta,
63+
a_index: 0f32,
64+
b_index: 0f32,
65+
a_dupe: None,
66+
b_dupe: None,
67+
}
68+
}
69+
70+
impl<I, J> Iterator for ZipStretch<I, J>
71+
where
72+
I: ExactSizeIterator,
73+
J: ExactSizeIterator,
74+
<I as Iterator>::Item: Clone,
75+
<J as Iterator>::Item: Clone,
76+
{
77+
type Item = (I::Item, J::Item);
78+
79+
fn next(&mut self) -> Option<Self::Item> {
80+
if self.a_index.fract() < self.a_delta {
81+
self.a_dupe = self.a.next();
82+
}
83+
self.a_index += self.a_delta;
84+
85+
if self.b_index.fract() < self.b_delta {
86+
self.b_dupe = self.b.next();
87+
}
88+
self.b_index += self.b_delta;
89+
90+
self.a_dupe.clone().zip(self.b_dupe.clone())
91+
}
92+
93+
fn size_hint(&self) -> (usize, Option<usize>) {
94+
size_hint::min(self.a.size_hint(), self.b.size_hint())
95+
}
96+
}
97+
98+
impl<I, J> ExactSizeIterator for ZipStretch<I, J>
99+
where
100+
I: ExactSizeIterator,
101+
J: ExactSizeIterator,
102+
<I as Iterator>::Item: Clone,
103+
<J as Iterator>::Item: Clone,
104+
{
105+
}

‎tests/zip.rs

+46
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,52 @@ use itertools::multizip;
22
use itertools::EitherOrBoth::{Both, Left, Right};
33
use itertools::Itertools;
44

5+
#[test]
6+
fn zip_squash() {
7+
let a: [i32; 0] = [];
8+
let b: [i32; 0] = [];
9+
let it = a.iter().copied().zip_squash(b.iter().copied());
10+
itertools::assert_equal(it, vec![]);
11+
12+
let a = [1, 2, 3, 4, 5, 6];
13+
let b: [i32; 0] = [];
14+
let it = a.iter().copied().zip_squash(b.iter().copied());
15+
itertools::assert_equal(it, vec![]);
16+
17+
let a: [i32; 0] = [];
18+
let b = [1, 2, 3, 7];
19+
let it = a.iter().copied().zip_squash(b.iter().copied());
20+
itertools::assert_equal(it, vec![]);
21+
22+
let a = [1, 2, 3, 4, 5, 6];
23+
let b = [1, 2, 3, 7];
24+
let it = a.iter().copied().zip_squash(b.iter().copied());
25+
itertools::assert_equal(it, vec![(1, 1), (2, 2), (4, 3), (5, 7)]);
26+
}
27+
28+
#[test]
29+
fn zip_stretch() {
30+
let a: [i32; 0] = [];
31+
let b: [i32; 0] = [];
32+
let it = a.iter().copied().zip_stretch(b.iter().copied());
33+
itertools::assert_equal(it, vec![]);
34+
35+
let a = [1, 2, 3, 4, 5, 6];
36+
let b: [i32; 0] = [];
37+
let it = a.iter().copied().zip_stretch(b.iter().copied());
38+
itertools::assert_equal(it, vec![]);
39+
40+
let a: [i32; 0] = [];
41+
let b = [1, 2, 3, 7];
42+
let it = a.iter().copied().zip_stretch(b.iter().copied());
43+
itertools::assert_equal(it, vec![]);
44+
45+
let a = [1, 2, 3, 4, 5, 6];
46+
let b = [1, 2, 3, 7];
47+
let it = a.iter().copied().zip_stretch(b.iter().copied());
48+
itertools::assert_equal(it, vec![(1, 1), (2, 1), (3, 2), (4, 3), (5, 3), (6, 7)]);
49+
}
50+
551
#[test]
652
fn zip_longest_fused() {
753
let a = [Some(1), None, Some(3), Some(4)];

0 commit comments

Comments
 (0)