1#![cfg(feature = "use_alloc")]
2
3use crate::size_hint;
4use crate::Itertools;
5
6use alloc::vec::Vec;
7
8#[derive(Clone)]
9/// An iterator adaptor that iterates over the cartesian product of
10/// multiple iterators of type `I`.
11///
12/// An iterator element type is `Vec<I>`.
13///
14/// See [`.multi_cartesian_product()`](crate::Itertools::multi_cartesian_product)
15/// for more information.
16#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
17pub struct MultiProduct<I>(Vec<MultiProductIter<I>>)
18where
19 I: Iterator + Clone,
20 I::Item: Clone;
21
22impl<I> std::fmt::Debug for MultiProduct<I>
23where
24 I: Iterator + Clone + std::fmt::Debug,
25 I::Item: Clone + std::fmt::Debug,
26{
27 debug_fmt_fields!(CoalesceBy, 0);
28}
29
30/// Create a new cartesian product iterator over an arbitrary number
31/// of iterators of the same type.
32///
33/// Iterator element is of type `Vec<H::Item::Item>`.
34pub fn multi_cartesian_product<H>(iters: H) -> MultiProduct<<H::Item as IntoIterator>::IntoIter>
35where
36 H: Iterator,
37 H::Item: IntoIterator,
38 <H::Item as IntoIterator>::IntoIter: Clone,
39 <H::Item as IntoIterator>::Item: Clone,
40{
41 MultiProduct(
42 itersimpl Iterator>
43 .map(|i: impl IntoIterator + Clone + Clone| MultiProductIter::new(i.into_iter()))
44 .collect(),
45 )
46}
47
48#[derive(Clone, Debug)]
49/// Holds the state of a single iterator within a `MultiProduct`.
50struct MultiProductIter<I>
51where
52 I: Iterator + Clone,
53 I::Item: Clone,
54{
55 cur: Option<I::Item>,
56 iter: I,
57 iter_orig: I,
58}
59
60/// Holds the current state during an iteration of a `MultiProduct`.
61#[derive(Debug)]
62enum MultiProductIterState {
63 StartOfIter,
64 MidIter { on_first_iter: bool },
65}
66
67impl<I> MultiProduct<I>
68where
69 I: Iterator + Clone,
70 I::Item: Clone,
71{
72 /// Iterates the rightmost iterator, then recursively iterates iterators
73 /// to the left if necessary.
74 ///
75 /// Returns true if the iteration succeeded, else false.
76 fn iterate_last(
77 multi_iters: &mut [MultiProductIter<I>],
78 mut state: MultiProductIterState,
79 ) -> bool {
80 use self::MultiProductIterState::*;
81
82 if let Some((last, rest)) = multi_iters.split_last_mut() {
83 let on_first_iter = match state {
84 StartOfIter => {
85 let on_first_iter = !last.in_progress();
86 state = MidIter { on_first_iter };
87 on_first_iter
88 }
89 MidIter { on_first_iter } => on_first_iter,
90 };
91
92 if !on_first_iter {
93 last.iterate();
94 }
95
96 if last.in_progress() {
97 true
98 } else if Self::iterate_last(rest, state) {
99 last.reset();
100 last.iterate();
101 // If iterator is None twice consecutively, then iterator is
102 // empty; whole product is empty.
103 last.in_progress()
104 } else {
105 false
106 }
107 } else {
108 // Reached end of iterator list. On initialisation, return true.
109 // At end of iteration (final iterator finishes), finish.
110 match state {
111 StartOfIter => false,
112 MidIter { on_first_iter } => on_first_iter,
113 }
114 }
115 }
116
117 /// Returns the unwrapped value of the next iteration.
118 fn curr_iterator(&self) -> Vec<I::Item> {
119 self.0
120 .iter()
121 .map(|multi_iter| multi_iter.cur.clone().unwrap())
122 .collect()
123 }
124
125 /// Returns true if iteration has started and has not yet finished; false
126 /// otherwise.
127 fn in_progress(&self) -> bool {
128 if let Some(last) = self.0.last() {
129 last.in_progress()
130 } else {
131 false
132 }
133 }
134}
135
136impl<I> MultiProductIter<I>
137where
138 I: Iterator + Clone,
139 I::Item: Clone,
140{
141 fn new(iter: I) -> Self {
142 Self {
143 cur: None,
144 iter: iter.clone(),
145 iter_orig: iter,
146 }
147 }
148
149 /// Iterate the managed iterator.
150 fn iterate(&mut self) {
151 self.cur = self.iter.next();
152 }
153
154 /// Reset the managed iterator.
155 fn reset(&mut self) {
156 self.iter = self.iter_orig.clone();
157 }
158
159 /// Returns true if the current iterator has been started and has not yet
160 /// finished; false otherwise.
161 fn in_progress(&self) -> bool {
162 self.cur.is_some()
163 }
164}
165
166impl<I> Iterator for MultiProduct<I>
167where
168 I: Iterator + Clone,
169 I::Item: Clone,
170{
171 type Item = Vec<I::Item>;
172
173 fn next(&mut self) -> Option<Self::Item> {
174 if Self::iterate_last(&mut self.0, MultiProductIterState::StartOfIter) {
175 Some(self.curr_iterator())
176 } else {
177 None
178 }
179 }
180
181 fn count(self) -> usize {
182 if self.0.is_empty() {
183 return 0;
184 }
185
186 if !self.in_progress() {
187 return self
188 .0
189 .into_iter()
190 .fold(1, |acc, multi_iter| acc * multi_iter.iter.count());
191 }
192
193 self.0.into_iter().fold(
194 0,
195 |acc,
196 MultiProductIter {
197 iter,
198 iter_orig,
199 cur: _,
200 }| {
201 let total_count = iter_orig.count();
202 let cur_count = iter.count();
203 acc * total_count + cur_count
204 },
205 )
206 }
207
208 fn size_hint(&self) -> (usize, Option<usize>) {
209 // Not ExactSizeIterator because size may be larger than usize
210 if self.0.is_empty() {
211 return (0, Some(0));
212 }
213
214 if !self.in_progress() {
215 return self.0.iter().fold((1, Some(1)), |acc, multi_iter| {
216 size_hint::mul(acc, multi_iter.iter.size_hint())
217 });
218 }
219
220 self.0.iter().fold(
221 (0, Some(0)),
222 |acc,
223 MultiProductIter {
224 iter,
225 iter_orig,
226 cur: _,
227 }| {
228 let cur_size = iter.size_hint();
229 let total_size = iter_orig.size_hint();
230 size_hint::add(size_hint::mul(acc, total_size), cur_size)
231 },
232 )
233 }
234
235 fn last(self) -> Option<Self::Item> {
236 let iter_count = self.0.len();
237
238 let lasts: Self::Item = self
239 .0
240 .into_iter()
241 .map(|multi_iter| multi_iter.iter.last())
242 .while_some()
243 .collect();
244
245 if lasts.len() == iter_count {
246 Some(lasts)
247 } else {
248 None
249 }
250 }
251}
252