itertools/adaptors/
multi_product.rs

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}
itertools/adaptors/multi_product.rs

itertools/adaptors/
multi_product.rs
