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