1#![cfg(feature = "use_alloc")]
2use Option::{self as State, None as ProductEnded, Some as ProductInProgress};
3use Option::{self as CurrentItems, None as NotYetPopulated, Some as Populated};
4
5use alloc::vec::Vec;
6
7use crate::size_hint;
8
9#[derive(Clone)]
10/// An iterator adaptor that iterates over the cartesian product of
11/// multiple iterators of type `I`.
12///
13/// An iterator element type is `Vec<I::Item>`.
14///
15/// See [`.multi_cartesian_product()`](crate::Itertools::multi_cartesian_product)
16/// for more information.
17#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
18pub struct MultiProduct<I>(State<MultiProductInner<I>>)
19where
20 I: Iterator + Clone,
21 I::Item: Clone;
22
23#[derive(Clone)]
24/// Internals for `MultiProduct`.
25struct MultiProductInner<I>
26where
27 I: Iterator + Clone,
28 I::Item: Clone,
29{
30 /// Holds the iterators.
31 iters: Vec<MultiProductIter<I>>,
32 /// Not populated at the beginning then it holds the current item of each iterator.
33 cur: CurrentItems<Vec<I::Item>>,
34}
35
36impl<I> std::fmt::Debug for MultiProduct<I>
37where
38 I: Iterator + Clone + std::fmt::Debug,
39 I::Item: Clone + std::fmt::Debug,
40{
41 debug_fmt_fields!(MultiProduct, 0);
42}
43
44impl<I> std::fmt::Debug for MultiProductInner<I>
45where
46 I: Iterator + Clone + std::fmt::Debug,
47 I::Item: Clone + std::fmt::Debug,
48{
49 debug_fmt_fields!(MultiProductInner, iters, cur);
50}
51
52/// Create a new cartesian product iterator over an arbitrary number
53/// of iterators of the same type.
54///
55/// Iterator element is of type `Vec<H::Item::Item>`.
56pub fn multi_cartesian_product<H>(iters: H) -> MultiProduct<<H::Item as IntoIterator>::IntoIter>
57where
58 H: Iterator,
59 H::Item: IntoIterator,
60 <H::Item as IntoIterator>::IntoIter: Clone,
61 <H::Item as IntoIterator>::Item: Clone,
62{
63 let inner: MultiProductInner<::IntoIter> = MultiProductInner {
64 iters: itersimpl Iterator>
65 .map(|i: impl IntoIterator + Clone + Clone| MultiProductIter::new(i.into_iter()))
66 .collect(),
67 cur: NotYetPopulated,
68 };
69 MultiProduct(ProductInProgress(inner))
70}
71
72#[derive(Clone, Debug)]
73/// Holds the state of a single iterator within a `MultiProduct`.
74struct MultiProductIter<I>
75where
76 I: Iterator + Clone,
77 I::Item: Clone,
78{
79 iter: I,
80 iter_orig: I,
81}
82
83impl<I> MultiProductIter<I>
84where
85 I: Iterator + Clone,
86 I::Item: Clone,
87{
88 fn new(iter: I) -> Self {
89 Self {
90 iter: iter.clone(),
91 iter_orig: iter,
92 }
93 }
94}
95
96impl<I> Iterator for MultiProduct<I>
97where
98 I: Iterator + Clone,
99 I::Item: Clone,
100{
101 type Item = Vec<I::Item>;
102
103 fn next(&mut self) -> Option<Self::Item> {
104 // This fuses the iterator.
105 let inner = self.0.as_mut()?;
106 match &mut inner.cur {
107 Populated(values) => {
108 debug_assert!(!inner.iters.is_empty());
109 // Find (from the right) a non-finished iterator and
110 // reset the finished ones encountered.
111 for (iter, item) in inner.iters.iter_mut().zip(values.iter_mut()).rev() {
112 if let Some(new) = iter.iter.next() {
113 *item = new;
114 return Some(values.clone());
115 } else {
116 iter.iter = iter.iter_orig.clone();
117 // `cur` is populated so the untouched `iter_orig` can not be empty.
118 *item = iter.iter.next().unwrap();
119 }
120 }
121 self.0 = ProductEnded;
122 None
123 }
124 // Only the first time.
125 NotYetPopulated => {
126 let next: Option<Vec<_>> = inner.iters.iter_mut().map(|i| i.iter.next()).collect();
127 if next.is_none() || inner.iters.is_empty() {
128 // This cartesian product had at most one item to generate and now ends.
129 self.0 = ProductEnded;
130 } else {
131 inner.cur.clone_from(&next);
132 }
133 next
134 }
135 }
136 }
137
138 fn count(self) -> usize {
139 match self.0 {
140 ProductEnded => 0,
141 // The iterator is fresh so the count is the product of the length of each iterator:
142 // - If one of them is empty, stop counting.
143 // - Less `count()` calls than the general case.
144 ProductInProgress(MultiProductInner {
145 iters,
146 cur: NotYetPopulated,
147 }) => iters
148 .into_iter()
149 .map(|iter| iter.iter_orig.count())
150 .try_fold(1, |product, count| {
151 if count == 0 {
152 None
153 } else {
154 Some(product * count)
155 }
156 })
157 .unwrap_or_default(),
158 // The general case.
159 ProductInProgress(MultiProductInner {
160 iters,
161 cur: Populated(_),
162 }) => iters.into_iter().fold(0, |mut acc, iter| {
163 if acc != 0 {
164 acc *= iter.iter_orig.count();
165 }
166 acc + iter.iter.count()
167 }),
168 }
169 }
170
171 fn size_hint(&self) -> (usize, Option<usize>) {
172 match &self.0 {
173 ProductEnded => (0, Some(0)),
174 ProductInProgress(MultiProductInner {
175 iters,
176 cur: NotYetPopulated,
177 }) => iters
178 .iter()
179 .map(|iter| iter.iter_orig.size_hint())
180 .fold((1, Some(1)), size_hint::mul),
181 ProductInProgress(MultiProductInner {
182 iters,
183 cur: Populated(_),
184 }) => {
185 if let [first, tail @ ..] = &iters[..] {
186 tail.iter().fold(first.iter.size_hint(), |mut sh, iter| {
187 sh = size_hint::mul(sh, iter.iter_orig.size_hint());
188 size_hint::add(sh, iter.iter.size_hint())
189 })
190 } else {
191 // Since it is populated, this cartesian product has started so `iters` is not empty.
192 unreachable!()
193 }
194 }
195 }
196 }
197
198 fn last(self) -> Option<Self::Item> {
199 let MultiProductInner { iters, cur } = self.0?;
200 // Collect the last item of each iterator of the product.
201 if let Populated(values) = cur {
202 let mut count = iters.len();
203 let last = iters
204 .into_iter()
205 .zip(values)
206 .map(|(i, value)| {
207 i.iter.last().unwrap_or_else(|| {
208 // The iterator is empty, use its current `value`.
209 count -= 1;
210 value
211 })
212 })
213 .collect();
214 if count == 0 {
215 // `values` was the last item.
216 None
217 } else {
218 Some(last)
219 }
220 } else {
221 iters.into_iter().map(|i| i.iter.last()).collect()
222 }
223 }
224}
225
226impl<I> std::iter::FusedIterator for MultiProduct<I>
227where
228 I: Iterator + Clone,
229 I::Item: Clone,
230{
231}
232