Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 014b99e

Browse files
iago-litoiago-lito
committedDec 8, 2024
🚧🚧 Refactor in terms of underlying pure-indices indices streaming iterator.
1 parent 3efe36c commit 014b99e

File tree

1 file changed

+248
-316
lines changed

1 file changed

+248
-316
lines changed
 

‎src/cartesian_power.rs

Lines changed: 248 additions & 316 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,206 @@
11
use alloc::vec::Vec;
22
use std::fmt;
33

4+
/// Pseudo-iterator owned by [`CartesianPower`],
5+
/// yielding underlying indices by references to itself,
6+
/// the [streaming iterator](https://docs.rs/streaming-iterator/latest/streaming_iterator/) way.
7+
#[derive(Debug, Clone)]
8+
pub struct Indices {
9+
// May be incremented by owner on first pass as long as exact value is unknown.
10+
base: usize,
11+
pow: u32,
12+
13+
// Indices just yielded. Length is 'pow'.
14+
// 0 0 .. 0 0 means that the first combination has been yielded.
15+
// 0 0 .. 0 1 means that the second combination has been yielded.
16+
// m m .. m m means that the last combination has just been yielded (m = base - 1).
17+
// b 0 .. 0 0 means that 'None' has just been yielded (b = base).
18+
// The latter is a special value marking the renewal of the iterator,
19+
// which can cycle again through another full round, ad libitum.
20+
values: Option<Vec<usize>>,
21+
}
22+
23+
impl Indices {
24+
pub fn new(base: usize, pow: u32) -> Self {
25+
Self {
26+
base,
27+
pow,
28+
values: None,
29+
}
30+
}
31+
pub fn next(&mut self) -> Option<&[usize]> {
32+
let Self { base, pow, values } = self;
33+
34+
match (base, pow, values) {
35+
// First iteration with degenerated 0th power.
36+
(_, 0, values @ None) => Some(values.insert(Vec::new())),
37+
38+
// Last degenerated 0th power iteration.
39+
// Use the Some<(empty)Vec> as a flag to alternate between yielding [] or None.
40+
(_, 0, values @ Some(_)) => {
41+
*values = None;
42+
None
43+
}
44+
45+
// Stable iteration in 0-base.
46+
(0, _, _) => None,
47+
48+
// First iteration in the general case.
49+
(_, pow, values @ None) => Some(values.insert(vec![0; *pow as usize])),
50+
51+
// Subsequent iteration in the general case.
52+
(&mut base, _, Some(values)) => {
53+
if values[0] == base {
54+
// Special marker that iteration can start over for a new round.
55+
values[0] = 0;
56+
return Some(values);
57+
}
58+
if inbounds_increment(values, base) {
59+
return Some(values);
60+
}
61+
// Iteration is over.
62+
// Mark a special index value to not fuse the iterator
63+
// and make it possible to cycle through all results again.
64+
values[0] = base;
65+
None
66+
}
67+
}
68+
}
69+
70+
pub fn nth(&mut self, n: usize) -> Option<&[usize]> {
71+
let Self { base, pow, values } = self;
72+
match (base, pow, values, n) {
73+
// First iteration with degenerated 0th power.
74+
(_, 0, values @ None, 0) => {
75+
// Same as .next()
76+
Some(values.insert(Vec::new()))
77+
}
78+
// Saturate.
79+
(_, 0, values @ Some(_), _) => {
80+
*values = None;
81+
None
82+
}
83+
// Stable iteration in 0-base.
84+
(0, _, _, _) => None,
85+
// First iteration in the general case.
86+
(&mut base, pow, values @ None, n) => {
87+
let values = values.insert(vec![0; *pow as usize]);
88+
if inbounds_increment_by(n, values, base) {
89+
return Some(values);
90+
}
91+
// Immediate saturation.
92+
values[0] = base;
93+
None
94+
}
95+
// Subsequent iteration in the general case.
96+
(&mut base, _, Some(values), n) => {
97+
let shift = if values[0] == base {
98+
// Start over for a new round (already counted then).
99+
values[0] = 0;
100+
0
101+
} else {
102+
1
103+
};
104+
if inbounds_increment_by(n + shift, values, base) {
105+
return Some(values);
106+
}
107+
// Immediate re-saturation.
108+
values[0] = base;
109+
None
110+
}
111+
}
112+
}
113+
114+
fn size_hint(&self) -> (usize, Option<usize>) {
115+
self.size_hint_with_base(self.base)
116+
}
117+
118+
fn size_hint_with_base(&self, base: usize) -> (usize, Option<usize>) {
119+
let Self {
120+
base: _,
121+
pow,
122+
values,
123+
} = self;
124+
125+
// The following case analysis matches implementation of `.next()`.
126+
match (base, pow, values) {
127+
// First iteration with degenerated 0th power.
128+
(_, 0, None) => (1, Some(1)),
129+
130+
// Last degenerated 0th power iteration | Stable iteration in 0-base.
131+
// Use the Some<(empty)Vec> as a flag to alternate between yielding [] or None.
132+
(0, _, _) | (_, 0, Some(_)) => (0, Some(0)),
133+
134+
// First iteration in the general case.
135+
(base, &pow, None) => {
136+
let c = base.checked_pow(pow);
137+
(c.unwrap_or(usize::MAX), c)
138+
}
139+
140+
// Subsequent iteration in the general case.
141+
(base, &pow, Some(values)) => {
142+
if values[0] == base {
143+
// Fresh re-start.
144+
let r = base.checked_pow(pow);
145+
return (r.unwrap_or(usize::MAX), r);
146+
}
147+
// Count what's left from current indices.
148+
// This is subtracting the current iteration number base^pow,
149+
// using the complement method.
150+
let calc = || -> Option<usize> {
151+
// (closure-wrap to ease interruption on overflow with ?-operator)
152+
let mut r = 0usize;
153+
for (&i, rank) in values.iter().rev().zip(0u32..) {
154+
let complement = base - 1 - i;
155+
let increment = complement.checked_mul(base.checked_pow(rank)?)?;
156+
r = r.checked_add(increment)?;
157+
}
158+
Some(r)
159+
};
160+
let Some(r) = calc() else {
161+
return (usize::MAX, None);
162+
};
163+
(r, Some(r))
164+
}
165+
}
166+
}
167+
}
168+
169+
/// Increment indices, returning false in case of overflow.
170+
fn inbounds_increment(indices: &mut [usize], base: usize) -> bool {
171+
for index in indices.iter_mut().rev() {
172+
*index += 1;
173+
if *index < base {
174+
return true;
175+
}
176+
*index = 0; // Wrap and increment left.
177+
}
178+
false
179+
}
180+
181+
/// Increment indices by n, returning false in case of (saturating) overflow.
182+
fn inbounds_increment_by(n: usize, indices: &mut [usize], base: usize) -> bool {
183+
let mut q = n;
184+
for index in indices.iter_mut().rev() {
185+
let s = *index + q;
186+
q = s / base;
187+
*index = s % base;
188+
if q == 0 {
189+
return true;
190+
}
191+
}
192+
// Saturation requires a second pass to reset all indices.
193+
for index in indices.iter_mut() {
194+
*index = 0;
195+
}
196+
false
197+
}
198+
4199
/// An adaptor iterating through all the ordered `n`-length lists of items
5200
/// yielded by the underlying iterator, including repetitions.
201+
/// Works by cloning the items into a fresh Vector on every `.next()`.
202+
/// If this is too much allocation for you,
203+
/// consider using the underlying lending [`Indices`] iterator instead.
6204
///
7205
/// See [`.cartesian_power()`](crate::Itertools::cartesian_power)
8206
/// for more information.
@@ -13,20 +211,11 @@ where
13211
I: Iterator,
14212
I::Item: Clone,
15213
{
16-
pow: u32,
17214
iter: Option<I>, // Inner iterator. Forget once consumed after 'base' iterations.
18-
items: Option<Vec<I::Item>>, // Fill from iter. Final length is 'base'.
19-
// None means that collection has not started yet.
20-
// Some(empty) means that collection would have started but pow = 0.
21-
22-
// Indices just yielded. Length is 'pow'.
23-
// 0 0 .. 0 0 means that the first combination has been yielded.
24-
// 0 0 .. 0 1 means that the second combination has been yielded.
25-
// m m .. m m means that the last combination has just been yielded (m = base - 1).
26-
// b 0 .. 0 0 means that 'None' has just been yielded (b = base).
27-
// The latter is a special value marking the renewal of the iterator,
28-
// which can cycle again through another full round, ad libitum.
29-
indices: Vec<usize>,
215+
items: Option<Vec<I::Item>>, // Fill from 'iter'. Final length is 'base'.
216+
// Keep track of the items to yield,
217+
// updating 'base' as the inner iterator is consumed.
218+
indices: Indices,
30219
}
31220

32221
/// Create a new `CartesianPower` from an iterator of clonables.
@@ -36,10 +225,9 @@ where
36225
I::Item: Clone,
37226
{
38227
CartesianPower {
39-
pow,
40228
iter: Some(iter),
41229
items: None,
42-
indices: Vec::new(),
230+
indices: Indices::new(0, pow),
43231
}
44232
}
45233

@@ -57,259 +245,53 @@ where
57245
/// valid within the collected items slice also returned.
58246
fn increment_indices(&mut self) -> Option<(&[usize], &[I::Item])> {
59247
let Self {
60-
pow,
61248
iter,
62249
items,
63250
indices,
64251
} = self;
65252

66-
let pow = *pow as usize;
67-
68-
// (weird 'items @' bindings circumvent NLL limitations, unneeded with polonius)
69-
match (pow, iter, &mut *items) {
70-
// First iteration with degenerated 0th power.
71-
(0, Some(_), items @ None) => {
72-
self.iter = None; // Forget about underlying iteration immediately.
73-
let empty = items.insert(Vec::new()); // Raise this value as a boolean flag.
74-
Some((indices, empty)) // Yield empty list.
75-
}
76-
77-
// Subsequent degenerated 0th power iteration.
78-
// Use the Some<(empty)Vec> as a flag to alternate between yielding [] or None.
79-
(0, None, items @ Some(_)) => {
80-
*items = None;
81-
None
82-
}
83-
(0, None, items @ None) => Some((indices, items.insert(Vec::new()))),
84-
85-
// First iteration in the general case.
86-
(pow, Some(it), items @ None) => {
87-
// Check whether there is at least one element in the iterator.
88-
if let Some(first) = it.next() {
89-
items // Collect it.
90-
.insert(Vec::with_capacity(it.size_hint().0))
91-
.push(first);
92-
// Prepare indices to be yielded.
93-
indices.reserve_exact(pow);
94-
for _ in 0..pow {
95-
indices.push(0);
96-
}
97-
Some((indices, items.as_ref().unwrap()))
98-
} else {
99-
// Degenerated iteration over an empty set:
100-
// 'base = 0', yet with non-null power.
101-
self.iter = None;
102-
None
103-
}
104-
}
105-
106-
// Stable iteration in the degenerated case 'base = 0'.
107-
(_, None, None) => None,
108-
109-
// Subsequent iteration in the general case.
110-
(pow, Some(it), Some(items)) => {
111-
// We are still unsure whether all items have been collected.
112-
// As a consequence, the exact value of 'base' is still uncertain,
113-
// but then we know that indices haven't started wrapping around yet.
114-
if let Some(next) = it.next() {
115-
items.push(next);
116-
indices[pow - 1] += 1;
117-
return Some((indices, items));
118-
}
119-
120-
// The item collected on previous call was the last one.
253+
if let Some(iter) = iter {
254+
let items = items.get_or_insert_with(|| Vec::with_capacity(iter.size_hint().0));
255+
if let Some(new) = iter.next() {
256+
indices.base += 1;
257+
items.push(new);
258+
} else {
121259
self.iter = None;
122-
let base = items.len(); // Definitive 'base' value.
123-
if base == 1 || pow == 1 {
124-
// Early end of singleton iteration.
125-
indices[0] = base; // Mark to cycle again on next iteration.
126-
return None;
127-
}
128-
129-
// First wrap around.
130-
indices[pow - 1] = 0;
131-
indices[pow - 2] = 1;
132-
Some((indices, items))
133-
}
134-
135-
// Subsequent iteration in the general case after all items have been collected.
136-
(_, None, Some(items)) => {
137-
let base = items.len();
138-
if indices[0] == base {
139-
// Special marker that iteration can start over for a new round.
140-
indices[0] = 0;
141-
return Some((indices, items));
142-
}
143-
// Keep yielding items list, incrementing indices rightmost first.
144-
if Self::inbounds_increment(indices, base) {
145-
return Some((indices, items));
146-
}
147-
// Iteration is over.
148-
// Mark a special index value to not fuse the iterator
149-
// and make it possible to cycle through all results again.
150-
indices[0] = base;
151-
None
152-
}
153-
}
154-
}
155-
156-
/// Increment indices, returning false in case of overflow.
157-
fn inbounds_increment(indices: &mut [usize], base: usize) -> bool {
158-
for index in indices.iter_mut().rev() {
159-
*index += 1;
160-
if *index < base {
161-
return true;
162-
}
163-
*index = 0; // Wrap and increment left.
164-
}
165-
false
166-
}
167-
168-
/// Increment indices by n, returning false in case of (saturating) overflow.
169-
fn inbounds_increment_by(n: usize, indices: &mut [usize], base: usize) -> bool {
170-
let mut q = n;
171-
for index in indices.iter_mut().rev() {
172-
let s = *index + q;
173-
q = s / base;
174-
*index = s % base;
175-
if q == 0 {
176-
return true;
177260
}
261+
indices.next().map(move |i| (i, items.as_slice()))
262+
} else if let Some(items) = items {
263+
indices.next().map(move |i| (i, items.as_slice()))
264+
} else {
265+
indices.next().map(move |i| (i, [].as_slice()))
178266
}
179-
// Saturation requires a second pass to reset all indices.
180-
for index in indices.iter_mut() {
181-
*index = 0;
182-
}
183-
false
184267
}
185268

186269
/// Same as [`increment_indices`], but does n increments at once.
187270
/// The iterator is cycling, but `.nth()` does not 'wrap'
188271
/// and 'saturates' to None instead.
189272
fn increment_indices_by_n(&mut self, n: usize) -> Option<(&[usize], &[I::Item])> {
190273
let Self {
191-
pow,
192274
iter,
193275
items,
194276
indices,
195277
} = self;
196278

197-
let pow = *pow as usize;
198-
199-
match (pow, iter, &mut *items, n) {
200-
// First iteration with degenerated 0th power.
201-
(0, Some(_), items @ None, 0) => {
202-
// Same as .next().
203-
self.iter = None;
204-
let empty = items.insert(Vec::new());
205-
Some((indices, empty))
206-
}
207-
(0, Some(_), None, _) => {
208-
// Saturate.
209-
self.iter = None;
210-
None
211-
}
212-
213-
// Subsequent degenerated 0th power iteration.
214-
// Same as `.next()`.
215-
(0, None, items @ None, 0) => Some((indices, items.insert(Vec::new()))),
216-
// Saturate.
217-
(0, None, items, _) => {
218-
*items = None;
219-
None
220-
}
221-
222-
// First iterations in the general case.
223-
// Possibly this will consume the entire underlying iterator,
224-
// but we need to consume to check.
225-
(pow, Some(it), items @ None, mut remaining) => {
226-
if let Some(first) = it.next() {
227-
// There is at least one element in the iterator, prepare collection + indices.
228-
let items = items.insert(Vec::with_capacity(it.size_hint().0));
229-
items.push(first);
230-
indices.reserve_exact(pow);
231-
for _ in 0..pow {
232-
indices.push(0);
233-
}
234-
// Collect more.
235-
loop {
236-
if remaining == 0 {
237-
// Stop before collection completion.
238-
indices[pow - 1] = n; // Hasn't wrapped yet.
239-
return Some((indices, items));
240-
}
241-
if let Some(next) = it.next() {
242-
items.push(next);
243-
remaining -= 1;
244-
continue;
245-
}
246-
// Collection completed, but we need to go further.
247-
self.iter = None;
248-
let base = items.len();
249-
if Self::inbounds_increment_by(n, indices, base) {
250-
return Some((indices, items));
251-
}
252-
// Immediate saturation.
253-
indices[0] = base;
254-
return None;
255-
}
279+
if let Some(iter) = iter {
280+
let items = items.get_or_insert_with(|| Vec::with_capacity(iter.size_hint().0));
281+
for _ in 0..n {
282+
if let Some(new) = iter.next() {
283+
indices.base += 1;
284+
items.push(new);
256285
} else {
257-
// Degenerated iteration over an empty set.
258286
self.iter = None;
259-
None
260-
}
261-
}
262-
263-
// Stable iteration in the degenerated case 'base = 0'.
264-
(_, None, None, _) => None,
265-
266-
// Subsequent iteration in the general case.
267-
// Again, immediate saturation is an option.
268-
(pow, Some(it), Some(items), mut remaining) => {
269-
if let Some(next) = it.next() {
270-
items.push(next);
271-
loop {
272-
if remaining == 0 {
273-
indices[pow - 1] += n + 1; // Hasn't wrapped yet.
274-
return Some((indices, items));
275-
}
276-
if let Some(next) = it.next() {
277-
items.push(next);
278-
remaining -= 1;
279-
continue;
280-
}
281-
break;
282-
}
283-
}
284-
// Collection completed.
285-
self.iter = None;
286-
let base = items.len();
287-
if Self::inbounds_increment_by(n + 1, indices, base) {
288-
return Some((indices, items));
287+
break;
289288
}
290-
// Saturate.
291-
indices[0] = base;
292-
None
293-
}
294-
295-
// Subsequent iteration in the general case
296-
// after all items have been collected.
297-
(_, None, Some(items), n) => {
298-
let base = items.len();
299-
let shift = if indices[0] == base {
300-
// Start over for a new round (already counted then).
301-
indices[0] = 0;
302-
0
303-
} else {
304-
1
305-
};
306-
if Self::inbounds_increment_by(n + shift, indices, base) {
307-
return Some((indices, items));
308-
}
309-
// Immediate re-saturation.
310-
indices[0] = base;
311-
None
312289
}
290+
indices.nth(n).map(move |i| (i, items.as_slice()))
291+
} else if let Some(items) = items {
292+
indices.nth(n).map(move |i| (i, items.as_slice()))
293+
} else {
294+
indices.nth(n).map(move |i| (i, [].as_slice()))
313295
}
314296
}
315297
}
@@ -343,75 +325,16 @@ where
343325
}
344326

345327
fn size_hint(&self) -> (usize, Option<usize>) {
346-
let Self {
347-
pow,
348-
iter,
349-
items,
350-
indices,
351-
} = self;
352-
353-
// The following case analysis matches implementation of `.next()`.
354-
#[allow(clippy::match_same_arms)]
355-
match (*pow, iter, items) {
356-
// First iteration with degenerated 0th power.
357-
(0, Some(_), None) => (1, Some(1)),
358-
359-
// Subsequent degenerated 0th power iteration.
360-
// Alternating for cycling behaviour.
361-
(0, None, Some(_)) => (0, Some(0)),
362-
(0, None, None) => (1, Some(1)),
363-
364-
// First iteration in the general case.
365-
(pow, Some(it), None) => {
366-
let (a, b) = it.size_hint();
367-
(
368-
a.checked_pow(pow).unwrap_or(usize::MAX),
369-
b.and_then(|b| b.checked_pow(pow)),
370-
)
371-
}
372-
373-
// Stable iteration in the degenerated case 'base = 0'.
374-
(_, None, None) => (0, Some(0)),
375-
376-
// Subsequent iteration in the general case.
377-
(pow, Some(it), Some(items)) => {
378-
let already = items.len();
379-
let minus_already = |total| total - already;
380-
let (a, b) = it.size_hint();
381-
(
382-
(a + already)
383-
.checked_pow(pow)
384-
.map_or(usize::MAX, minus_already),
385-
b.and_then(|b| (b + already).checked_pow(pow).map(minus_already)),
386-
)
387-
}
388-
389-
// Subsequent iteration in the general case after all items have been collected.
390-
(pow, None, Some(items)) => {
391-
let base = items.len();
392-
if indices[0] == base {
393-
// Fresh re-start.
394-
let r = base.checked_pow(pow);
395-
return (r.unwrap_or(usize::MAX), r);
396-
}
397-
// Count what's left from current indices.
398-
// This is subtracting the current iteration number base^pow,
399-
// using the complement method.
400-
let calc = || -> Option<usize> {
401-
// (closure-wrap to ease interruption on overflow with ?-operator)
402-
let mut r = 0usize;
403-
for (&i, rank) in indices.iter().rev().zip(0u32..) {
404-
let complement = base - 1 - i;
405-
let increment = complement.checked_mul(base.checked_pow(rank)?)?;
406-
r = r.checked_add(increment)?;
407-
}
408-
Some(r)
409-
};
410-
let Some(r) = calc() else {
411-
return (usize::MAX, None);
412-
};
413-
(r, Some(r))
414-
}
328+
let Self { iter, indices, .. } = self;
329+
// Forward low/high hints from underlying iterator
330+
// to indices iterators.
331+
if let Some(iter) = iter {
332+
let (a, b) = iter.size_hint();
333+
let a = indices.size_hint_with_base(a).0;
334+
let b = b.and_then(|b| indices.size_hint_with_base(b).1);
335+
(a, b)
336+
} else {
337+
indices.size_hint()
415338
}
416339
}
417340

@@ -428,13 +351,11 @@ where
428351
{
429352
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
430353
let Self {
431-
pow,
432354
iter,
433355
items,
434356
indices,
435357
} = self;
436358
f.debug_struct("CartesianPower")
437-
.field("pow", pow)
438359
.field("iter", &iter.is_some())
439360
.field("items", items)
440361
.field("indices", indices)
@@ -446,8 +367,19 @@ where
446367
#[cfg(test)]
447368
mod tests {
448369

449-
use crate::Itertools;
370+
use crate::{cartesian_power::Indices, Itertools};
450371

372+
#[test]
373+
fn indices() {
374+
let mut it = Indices::new(3, 2);
375+
for i in 0..30 {
376+
println!("{i}: {:?}", it.next());
377+
}
378+
for i in 0..30 {
379+
println!("{i}: {:?}", it.nth(2));
380+
}
381+
panic!("STOP HERE");
382+
}
451383
#[test]
452384
fn basic() {
453385
fn check(origin: &str, pow: u32, expected: &[&str]) {

0 commit comments

Comments
 (0)
Please sign in to comment.