cartesian/
lib.rs

1#![no_std]
2
3//! Provides a quality of life macro `cartesian!` to simplify certain loops.
4//!
5//! The macro takes up to 26 iterators as arguments and creates the cartesian product iterator over
6//! all input iterators, kind of like nested for loops.
7//!
8//! It behaves the same as nested for loops and brings the advantage of
9//! being more compact, and simplifies breaking and continuing.
10//!
11//! # Examples
12//! ```
13//! use cartesian::cartesian;
14//!
15//! let mut volume_grid = vec![vec![vec![0; 10]; 10]; 10];
16//! for (x, y, z) in cartesian!(0..10, 0..10, 0..10) {
17//!     volume_grid[x][y][z] = x * y + z;
18//! }
19//! ```
20//!
21//! ```
22//! # use cartesian::cartesian;
23//! let (m, n, p) = (3, 3, 1);
24//!
25//! let mut id = vec![vec![0; n]; m];
26//! for (i, j) in cartesian!(0..m, 0..n) {
27//!     id[i][j] = (i == j) as u32;
28//! }
29//!
30//! let col_vec = vec![vec![1], vec![2], vec![4]];
31//!
32//! let mut res = vec![vec![0; p]; m];
33//!
34//! for (i, j, k) in cartesian!(0..m, 0..n, 0..p) {
35//!     res[i][k] += id[i][j] * col_vec[j][k];
36//! }
37//!
38//! assert_eq!(col_vec, res);
39//! ```
40
41/// Helper trait implemented for all tuples up to 26 elements to prepend a value to produce a longer tuple
42///
43/// The implementation is adapted from [this stackoverflow answer](https://stackoverflow.com/a/57454888).
44pub trait TuplePrepend<T> {
45    type ResultType;
46    fn prepend(self, t: T) -> Self::ResultType;
47}
48
49impl<T> TuplePrepend<T> for () {
50    type ResultType = (T,);
51    #[inline]
52    fn prepend(self, t: T) -> Self::ResultType {
53        (t,)
54    }
55}
56
57macro_rules! _impl_tuple_prepend {
58    ( ()  ) => {};
59    ( ($t:ident $( $typ:ident)* )  ) => {
60        impl<$t, $($typ,)* TT> TuplePrepend<TT> for ($t, $($typ,)*) {
61            type ResultType = (TT, $t, $($typ),*);
62            #[inline]
63            fn prepend(self, t: TT) -> Self::ResultType {
64                #[allow(non_snake_case)]
65                let ($t, $($typ,)*) = self;
66                (t, $t, $($typ,)*)
67
68            }
69
70        }
71        _impl_tuple_prepend!(($($typ)*));
72
73    }
74
75}
76_impl_tuple_prepend!((
77    A B C D E F G H I J K L M N O P Q R S T U V W X Y Z
78));
79
80/// The macro this is all about.
81#[macro_export]
82macro_rules! cartesian {
83    ($head:expr $(,)?) => {
84        $head.into_iter()
85    };
86
87    // We need this base case to ensure that
88    // 1. `.map(|(head, tail)| tail.prepend(head))` is only called when `tail` is a tuple type.
89    // 2. `.into_iter()` is called exactly once for each macro argument.
90    ($head:expr, $tail:expr $(,)?) => {
91        cartesian!(@ $head.into_iter(), $tail.into_iter())
92    };
93
94    // Expression                                              | Type
95    // --------------------------------------------------------+----------------------------------------
96    // $head.into_iter()                                       | impl Iterator<Item = A>
97    // cartesian!($($tail),+)                                  | impl Iterator<Item = (B, C, ...)>
98    // cartesian!(@ $head.into_iter(), cartesian!($($tail),+)) | impl Iterator<Item = (A, (B, C, ...))>
99    // cartesian!(...).map(...)                                | impl Iterator<Item = (A, B, C, ...)>
100    ($head:expr $(, $tail:expr)+ $(,)?) => {
101        cartesian!(@ $head.into_iter(), cartesian!($($tail),+)).map(
102            |(head, tail)| $crate::TuplePrepend::prepend(tail, head)
103        )
104    };
105
106    (@ $head:expr, $tail:expr $(,)?) => {
107        $head.flat_map(|h| $tail.map(move |t| (h, t)))
108    };
109}
110
111#[cfg(test)]
112extern crate alloc;
113#[cfg(test)]
114use alloc::{format, string::String, vec};
115
116#[test]
117fn two_combination() {
118    let mut acc = String::new();
119
120    for (a, b) in cartesian!(0..2, "xy".chars()) {
121        acc += &format!("{}{} ", a, b);
122    }
123
124    assert_eq!(acc, "0x 0y 1x 1y ");
125}
126
127#[test]
128fn binary_numbers() {
129    let mut acc = String::new();
130
131    let range = 0..2;
132    let vec = vec![0, 1];
133    let string = vec![String::from("0"), String::from("1")];
134
135    for (a, b, c) in cartesian!(range, vec.iter(), string.iter()) {
136        acc += &format!("{}{}{} ", a, b, c);
137    }
138
139    assert_eq!(acc, "000 001 010 011 100 101 110 111 ");
140}
141
142#[test]
143fn trailing_commas() {
144    let mut acc = String::new();
145
146    for a in cartesian!(
147        0..1,
148    ) {
149        acc += &format!("{} ", a);
150    }
151
152    for (a, b) in cartesian!(
153        0..2,
154        0..2,
155    ) {
156        acc += &format!("{}{} ", a, b);
157    }
158
159    for (a, b, c) in cartesian!(
160        0..2,
161        0..2,
162        0..2,
163    ) {
164        acc += &format!("{}{}{} ", a, b, c);
165    }
166
167    assert_eq!(acc, "0 00 01 10 11 000 001 010 011 100 101 110 111 ");
168}
169
170#[test]
171fn by_reference() {
172    let mut acc = String::new();
173
174    let outer = vec![String::from("a"), String::from("b")];
175    let inner = vec![String::from("0"), String::from("1")];
176
177    for (a, b) in cartesian!(&outer, &inner) {
178        acc += &format!("{}{} ", a, b);
179    }
180
181    assert_eq!(acc, "a0 a1 b0 b1 ");
182}
cartesian/lib.rs

cartesian/
lib.rs
