Skip to content

Commit 8957f70

Browse files
iago-litoiago-lito
committed
🚧 Implementing cartesian power.
1 parent 4777762 commit 8957f70

File tree

2 files changed

+376
-0
lines changed

2 files changed

+376
-0
lines changed

‎src/cartesian_power.rs

+356
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,356 @@
1+
use alloc::vec::Vec;
2+
use std::fmt;
3+
use std::iter::FusedIterator;
4+
5+
/// An adaptor iterating through all the ordered `n`-length lists of items
6+
/// yielded by the underlying iterator, including repetitions.
7+
///
8+
/// See [`.cartesian_power()`](crate::Itertools::cartesian_power)
9+
/// for more information.
10+
#[derive(Clone)]
11+
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
12+
pub struct CartesianPower<I>
13+
where
14+
I: Iterator,
15+
I::Item: Clone,
16+
{
17+
pow: usize,
18+
iter: Option<I>, // Inner iterator. Forget once consumed after 'base' iterations.
19+
items: Vec<I::Item>, // Fill from iter. Clear once adaptor is exhausted. Final length is 'base'.
20+
indices: Vec<usize>, // Indices just yielded. Clear once adaptor is exhausted. Length is 'pow'.
21+
}
22+
23+
/// Create a new `CartesianPower` from an iterator of clonables.
24+
pub fn cartesian_power<I>(iter: I, pow: usize) -> CartesianPower<I>
25+
where
26+
I: Iterator,
27+
I::Item: Clone,
28+
{
29+
CartesianPower {
30+
pow,
31+
iter: Some(iter),
32+
items: Vec::new(),
33+
indices: Vec::new(),
34+
}
35+
}
36+
37+
impl<I> CartesianPower<I>
38+
where
39+
I: Iterator,
40+
I::Item: Clone,
41+
{
42+
/// Increments internal indices to advance to the next list to be yielded.
43+
/// This collects new items from the underlying iterator
44+
/// if they were not all already collected.
45+
///
46+
/// Returns None if we've run out of possible lists,
47+
/// otherwise return refs to the indices to yield next,
48+
/// valid within the collected items slice also returned.
49+
fn increment_indices(&mut self) -> Option<(&[usize], &[I::Item])> {
50+
let Self {
51+
pow,
52+
iter,
53+
items,
54+
indices,
55+
} = self;
56+
match (*pow, iter, items.len()) {
57+
// Final stable state: underlying iterator and items forgotten.
58+
(_, None, 0) => None,
59+
60+
// Degenerated 0th power iteration.
61+
(0, Some(_), _) => {
62+
self.iter = None; // Forget without even consuming.
63+
Some((indices, items))
64+
}
65+
66+
(pow, Some(it), 0) => {
67+
// Check whether there is at least one element in the iterator.
68+
if let Some(first) = it.next() {
69+
// Allocate buffer to hold items about to be yielded.
70+
items.reserve_exact(it.size_hint().0);
71+
items.push(first);
72+
// Same for indices to be yielded.
73+
indices.reserve_exact(pow);
74+
for _ in 0..pow {
75+
indices.push(0);
76+
}
77+
return Some((indices, items));
78+
}
79+
// Degenerated iteration over an empty set, yet with non-null power.
80+
self.iter = None;
81+
None
82+
}
83+
84+
(pow, Some(it), base) => {
85+
// We are still unsure whether all items have been collected.
86+
// As a consequence, 'base' is still uncertain,
87+
// but then we know that indices haven't started wrapping around yet.
88+
if let Some(next) = it.next() {
89+
items.push(next);
90+
indices[pow - 1] += 1;
91+
return Some((indices, items));
92+
}
93+
94+
// All items have just been collected.
95+
self.iter = None;
96+
if base == 1 || pow == 1 {
97+
// End of iteration.
98+
items.clear();
99+
indices.clear();
100+
return None;
101+
}
102+
103+
// First wrap around.
104+
indices[pow - 1] = 0;
105+
indices[pow - 2] += 1;
106+
Some((indices, items))
107+
}
108+
109+
(_, None, b) => {
110+
// Keep yielding items list, incrementing indices rightmost first.
111+
for index in indices.iter_mut().rev() {
112+
*index += 1;
113+
if *index < b {
114+
return Some((indices, items));
115+
}
116+
*index = 0; // Wrap and increment left.
117+
}
118+
items.clear();
119+
indices.clear();
120+
None
121+
}
122+
}
123+
}
124+
125+
/// Same as [`increment_indices`], but does n increments at once.
126+
fn increment_indices_by_n(&mut self, n: usize) -> Option<(&[usize], &[I::Item])> {
127+
todo!()
128+
}
129+
}
130+
131+
impl<I> Iterator for CartesianPower<I>
132+
where
133+
I: Iterator,
134+
I::Item: Clone,
135+
{
136+
type Item = Vec<I::Item>;
137+
138+
fn next(&mut self) -> Option<Self::Item> {
139+
// If anything to yield,
140+
// clone the correct 'pow' instances of collected items
141+
// into a freshly allocated vector.
142+
self.increment_indices().map(|(indices, items)| {
143+
indices
144+
.iter()
145+
.map(|&i| items[i].clone())
146+
.collect::<Vec<_>>()
147+
})
148+
}
149+
150+
fn nth(&mut self, n: usize) -> Option<Self::Item> {
151+
self.increment_indices_by_n(n).map(|(indices, items)| {
152+
indices
153+
.iter()
154+
.map(|&i| items[i].clone())
155+
.collect::<Vec<_>>()
156+
})
157+
}
158+
159+
fn size_hint(&self) -> (usize, Option<usize>) {
160+
todo!()
161+
}
162+
163+
fn count(self) -> usize {
164+
todo!()
165+
}
166+
}
167+
168+
// Elide underlying iterator from the debug display.
169+
impl<I> fmt::Debug for CartesianPower<I>
170+
where
171+
I: Iterator + fmt::Debug,
172+
I::Item: fmt::Debug + Clone,
173+
{
174+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
175+
let Self {
176+
pow,
177+
iter,
178+
items,
179+
indices,
180+
} = self;
181+
f.debug_struct("CartesianPower")
182+
.field("pow", pow)
183+
.field("iter", &iter.is_some())
184+
.field("items", items)
185+
.field("indices", indices)
186+
.finish()
187+
}
188+
}
189+
190+
impl<I> FusedIterator for CartesianPower<I>
191+
where
192+
I: Iterator,
193+
I::Item: Clone,
194+
{
195+
}
196+
197+
#[cfg(test)]
198+
mod tests {
199+
//! Use chars and string to ease testing of every yielded iterator values.
200+
201+
use super::CartesianPower;
202+
use crate::Itertools;
203+
use core::str::Chars;
204+
205+
fn check_fused(mut exhausted_it: CartesianPower<Chars>, context: String) {
206+
for i in 0..100 {
207+
let act = exhausted_it.next();
208+
assert!(
209+
act.is_none(),
210+
"Iteration {} after expected exhaustion of {} \
211+
yielded {:?} instead of None. ",
212+
i,
213+
context,
214+
act,
215+
);
216+
}
217+
}
218+
219+
#[test]
220+
fn basic() {
221+
fn check(origin: &str, pow: usize, expected: &[&str]) {
222+
let mut it = origin.chars().cartesian_power(pow);
223+
let mut i = 0;
224+
for exp in expected {
225+
let act = it.next();
226+
if act != Some(exp.chars().collect()) {
227+
panic!(
228+
"Failed iteration {} for {:?}^{}. \
229+
Expected {:?}, got {:?} instead.",
230+
i, origin, pow, exp, act,
231+
);
232+
}
233+
i += 1;
234+
}
235+
check_fused(it, format!("iteration {} or {:?}^{}", i, origin, pow));
236+
}
237+
238+
// Empty underlying iterator.
239+
check("", 0, &[""]);
240+
check("", 1, &[]);
241+
check("", 2, &[]);
242+
check("", 3, &[]);
243+
244+
// Singleton underlying iterator.
245+
check("a", 0, &[""]);
246+
check("a", 1, &["a"]);
247+
check("a", 2, &["aa"]);
248+
check("a", 3, &["aaa"]);
249+
250+
// Underlying pair.
251+
check("ab", 0, &[""]);
252+
check("ab", 1, &["a", "b"]);
253+
check("ab", 2, &["aa", "ab", "ba", "bb"]);
254+
check(
255+
"ab",
256+
3,
257+
&["aaa", "aab", "aba", "abb", "baa", "bab", "bba", "bbb"],
258+
);
259+
260+
// Underlying triplet.
261+
check("abc", 0, &[""]);
262+
check("abc", 1, &["a", "b", "c"]);
263+
check(
264+
"abc",
265+
2,
266+
&["aa", "ab", "ac", "ba", "bb", "bc", "ca", "cb", "cc"],
267+
);
268+
check(
269+
"abc",
270+
3,
271+
&[
272+
"aaa", "aab", "aac", "aba", "abb", "abc", "aca", "acb", "acc", "baa", "bab", "bac",
273+
"bba", "bbb", "bbc", "bca", "bcb", "bcc", "caa", "cab", "cac", "cba", "cbb", "cbc",
274+
"cca", "ccb", "ccc",
275+
],
276+
);
277+
}
278+
279+
#[test]
280+
fn nth() {
281+
fn check(origin: &str, pow: usize, expected: &[(usize, Option<&str>)]) {
282+
let mut it = origin.chars().cartesian_power(pow);
283+
let mut total_n = Vec::new();
284+
for &(n, exp) in expected {
285+
let act = it.nth(n);
286+
total_n.push(n);
287+
if act != exp.map(|s| s.chars().collect::<Vec<_>>()) {
288+
panic!(
289+
"Failed nth({}) iteration for {:?}^{}. \
290+
Expected {:?}, got {:?} instead.",
291+
total_n
292+
.iter()
293+
.map(ToString::to_string)
294+
.collect::<Vec<_>>()
295+
.join(", "),
296+
origin,
297+
pow,
298+
exp,
299+
act,
300+
);
301+
}
302+
}
303+
check_fused(
304+
it,
305+
format!(
306+
"nth({}) iteration of {:?}^{}",
307+
total_n
308+
.iter()
309+
.map(ToString::to_string)
310+
.collect::<Vec<_>>()
311+
.join(", "),
312+
origin,
313+
pow
314+
),
315+
);
316+
}
317+
318+
// HERE: make it work with the new implementation.
319+
320+
// Check degenerated cases.
321+
check("", 0, &[(0, Some("")), (0, None)]);
322+
check("", 0, &[(0, Some("")), (1, None)]);
323+
check("", 0, &[(0, Some("")), (2, None)]);
324+
check("", 0, &[(1, None), (0, None)]);
325+
check("", 0, &[(1, None), (1, None)]);
326+
check("", 0, &[(1, None), (2, None)]);
327+
check("", 0, &[(2, None), (0, None)]);
328+
check("", 0, &[(2, None), (1, None)]);
329+
check("", 0, &[(2, None), (2, None)]);
330+
331+
check("a", 0, &[(0, Some("")), (0, None)]);
332+
check("a", 0, &[(0, Some("")), (1, None)]);
333+
check("a", 0, &[(0, Some("")), (2, None)]);
334+
check("a", 0, &[(1, None), (0, None)]);
335+
check("a", 0, &[(1, None), (1, None)]);
336+
check("a", 0, &[(1, None), (2, None)]);
337+
check("a", 0, &[(2, None), (0, None)]);
338+
check("a", 0, &[(2, None), (1, None)]);
339+
check("a", 0, &[(2, None), (2, None)]);
340+
341+
// Unit power.
342+
check("a", 1, &[(0, Some("a")), (0, None)]);
343+
check("a", 1, &[(0, Some("a")), (1, None)]);
344+
check("a", 1, &[(0, Some("a")), (2, None)]);
345+
check("a", 1, &[(1, None), (0, None)]);
346+
check("a", 1, &[(1, None), (1, None)]);
347+
check("a", 1, &[(1, None), (2, None)]);
348+
check("a", 1, &[(2, None), (0, None)]);
349+
check("a", 1, &[(2, None), (1, None)]);
350+
check("a", 1, &[(2, None), (2, None)]);
351+
352+
check("ab", 1, &[(0, Some("a")), (0, Some("b")), (0, None)]);
353+
check("ab", 1, &[(1, Some("b")), (0, None), (0, None)]);
354+
check("ab", 1, &[(2, None), (0, None), (0, None)]);
355+
}
356+
}

‎src/lib.rs

+20
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ extern crate alloc;
5757
#[cfg(feature = "use_alloc")]
5858
use alloc::{collections::VecDeque, string::String, vec::Vec};
5959

60+
use cartesian_power::CartesianPower;
6061
pub use either::Either;
6162

6263
use core::borrow::Borrow;
@@ -175,6 +176,8 @@ pub mod free;
175176
#[doc(inline)]
176177
pub use crate::free::*;
177178
#[cfg(feature = "use_alloc")]
179+
mod cartesian_power;
180+
#[cfg(feature = "use_alloc")]
178181
mod combinations;
179182
#[cfg(feature = "use_alloc")]
180183
mod combinations_with_replacement;
@@ -1729,6 +1732,23 @@ pub trait Itertools: Iterator {
17291732
combinations_with_replacement::combinations_with_replacement(self, k)
17301733
}
17311734

1735+
/// Returns an iterator yielding the successive elements
1736+
/// of the cartesian power of the set described by the original iterator.
1737+
///
1738+
/// ```
1739+
/// use itertools::Itertools;
1740+
///
1741+
/// TODO: illustrative example.
1742+
/// ```
1743+
#[cfg(feature = "use_alloc")]
1744+
fn cartesian_power(self, pow: usize) -> CartesianPower<Self>
1745+
where
1746+
Self: Sized,
1747+
Self::Item: Clone,
1748+
{
1749+
cartesian_power::cartesian_power(self, pow)
1750+
}
1751+
17321752
/// Return an iterator adaptor that iterates over all k-permutations of the
17331753
/// elements from an iterator.
17341754
///

0 commit comments

Comments
 (0)