itermore/adaptors/
cartesian_product.rs

1use core::iter::FusedIterator;
2
3/// An extension trait that provides the [`cartesian_product`] method for
4/// iterators.
5///
6/// [`cartesian_product`]: IterCartesianProduct::cartesian_product
7pub trait IterCartesianProduct: Iterator {
8    /// Returns an iterator adaptor that iterates over the cartesian product of
9    /// the element sets of two iterators `self` and `other.into_iter()`.
10    ///
11    /// # Examples
12    ///
13    /// ```
14    /// use itermore::IterCartesianProduct;
15    ///
16    /// let v = Vec::from_iter((0..3).cartesian_product("αβ".chars()));
17    /// assert_eq!(v, [(0, 'α'), (0, 'β'), (1, 'α'), (1, 'β'), (2, 'α'), (2, 'β')]);
18    /// ```
19    fn cartesian_product<J>(self, other: J) -> CartesianProduct<Self, J::IntoIter>
20    where
21        Self: Sized,
22        Self::Item: Clone,
23        J: IntoIterator,
24        J::IntoIter: Clone,
25    {
26        CartesianProduct::new(self, other.into_iter())
27    }
28}
29
30impl<I: ?Sized> IterCartesianProduct for I where I: Iterator {}
31
32/// An iterator over the cartesian product of the element sets of two iterators
33/// `I` and `J`.
34///
35/// This struct is created by the [`cartesian_product`] method on iterators. See
36/// its documentation for more.
37///
38/// [`cartesian_product`]: IterCartesianProduct::cartesian_product
39#[derive(Debug, Clone)]
40#[must_use = "iterators are lazy and do nothing unless consumed"]
41pub struct CartesianProduct<I, J>
42where
43    I: Iterator,
44{
45    a: I,
46    b: J,
47    a_item: Option<I::Item>,
48    b_curr: J,
49}
50
51impl<I, J> CartesianProduct<I, J>
52where
53    I: Iterator,
54    J: Iterator + Clone,
55{
56    fn new(mut a: I, b: J) -> Self {
57        CartesianProduct {
58            a_item: a.next(),
59            a,
60            b_curr: b.clone(),
61            b,
62        }
63    }
64}
65
66impl<I, J> Iterator for CartesianProduct<I, J>
67where
68    I: Iterator,
69    J: Iterator + Clone,
70    I::Item: Clone,
71{
72    type Item = (I::Item, J::Item);
73
74    fn next(&mut self) -> Option<Self::Item> {
75        let b_item = match self.b_curr.next() {
76            Some(b_item) => b_item,
77            None => {
78                self.b_curr = self.b.clone();
79                let b_item = self.b_curr.next()?;
80                self.a_item = self.a.next();
81                b_item
82            }
83        };
84        self.a_item.as_ref().map(|a| (a.clone(), b_item))
85    }
86}
87
88impl<I, J> FusedIterator for CartesianProduct<I, J>
89where
90    I: FusedIterator,
91    J: FusedIterator + Clone,
92    I::Item: Clone,
93{
94}
95
96////////////////////////////////////////////////////////////////////////////////
97// Macro
98////////////////////////////////////////////////////////////////////////////////
99
100/// Returns an iterator over the cartesian product of the element sets of
101/// multiple iterators (up to 12).
102///
103/// This is essentially the equivalent of calling [`cartesian_product`] multiple
104/// times and "flattening" each item e.g. `((A, B), C)` to `(A, B, C)`.
105///
106/// # Examples
107///
108/// ```
109/// use itermore::{cartesian_product, IterCartesianProduct};
110///
111/// // With macro
112/// let i = cartesian_product!(0..3, "αβ".chars(), [-1, 0, 1]);
113///
114/// // Without macro
115/// let j = (0..3)
116///     .cartesian_product("αβ".chars())
117///     .cartesian_product([-1, 0, 1])
118///     .map(|((a, b), c)| (a, b, c));
119///
120/// assert_eq!(Vec::from_iter(i), Vec::from_iter(j));
121/// ```
122///
123/// Iterate over the 3D coordinates of a 10 x 10 x 10 cube.
124///
125/// ```
126/// use itermore::cartesian_product;
127///
128/// // With macro
129/// for (i, j, k) in cartesian_product!(0..10, 0..10, 0..10) {
130///     // ...
131/// }
132///
133/// // Without macro
134/// for i in 0..10 {
135///     for j in 0..10 {
136///         for k in 0..10 {
137///             // ...
138///         }
139///     }
140/// }
141/// ```
142///
143/// [`cartesian_product`]: IterCartesianProduct::cartesian_product
144#[macro_export]
145macro_rules! cartesian_product {
146    ($I:expr $(,)?) => {
147        $crate::core::iter::IntoIterator::into_iter($I)
148    };
149
150    ($I:expr, $J:expr $(,)?) => {
151        $crate::IterCartesianProduct::cartesian_product(
152            $crate::cartesian_product!($I),
153            $crate::cartesian_product!($J),
154        )
155    };
156
157    ($I:expr, $J:expr, $($K:expr),+ $(,)?) => {{
158        $crate::cartesian_product!($crate::cartesian_product!($I, $J), $($K),+)
159            .map($crate::flatten_tuple)
160    }};
161}
itermore/adaptors/cartesian_product.rs

itermore/adaptors/
cartesian_product.rs
