rustcomp/
lib.rs

1#![warn(clippy::all, clippy::pedantic)]
2/*!
3Adds comprehensions to Rust. This is achieved through a functional macro,
4[`rcomp!`], that does all the heavy lifting for you.
5
6# Basic Usage
7
8The core idea is simple: provide an easy and concise way to flatten, filter,
9map, and collect iterators. For a full breakdown of the syntax, see
10the docs for the [`rcomp!`] macro. For now, consider this simple example:
11
12```rust
13# use rustcomp::rcomp;
14let v = rcomp![Vec<_>; for x in 0..10 => x];
15let it = (0..10).collect::<Vec<_>>(); // all examples show an equivalent iterator
16assert_eq!(v, it);
17```
18
19This will make a `Vec<i32>` with the numbers 0 through 9... not very useful,
20is it? Let's add a guard to filter out the odd numbers:
21
22```rust
23# use rustcomp::rcomp;
24let v = rcomp![Vec<_>; for x in 0..10 => x, if x % 2 == 0];
25let it = (0..10).filter(|x| x % 2 == 0).collect::<Vec<_>>();
26assert_eq!(v, it);
27```
28
29Now we're getting somewhere! You can also map the values, so let's double
30them for fun:
31
32```rust
33# use rustcomp::rcomp;
34let v = rcomp![Vec<_>; for x in 0..10 => x * 2, if x % 2 == 0];
35let it = (0..10)
36    .filter(|x| x % 2 == 0)
37    .map(|x| x * 2)
38    .collect::<Vec<_>>();
39assert_eq!(v, it);
40```
41
42Notice how the `map` call comes _after_ the `filter` call in the iterator example.
43This is also how the comprehension works: the guard applies to the _input_ value,
44not the output value.
45
46Speaking of iterators, if you don't want to collect the results into a container,
47you can get the iterator directly by omitting the collection type:
48
49```rust
50# use rustcomp::rcomp;
51// now we have to collect the iterator ourselves
52let v = rcomp![for x in 0..10 => x].collect::<Vec<_>>();
53// equivalent to:
54let vv = rcomp![Vec<_>; for x in 0..10 => x];
55# let it = (0..10)
56#     .collect::<Vec<_>>();
57# assert_eq!(v, vv);
58# assert_eq!(v, it);
59```
60
61# Destructuring
62
63Comprehensions also support destructuring. For example, tuples:
64
65```rust
66# use rustcomp::rcomp;
67let pairs = vec![(1, 2), (3, 4), (5, 6)];
68let v = rcomp![Vec<_>; for (x, y) in &pairs => x + y];
69let it = pairs.into_iter().map(|(x, y)| x + y).collect::<Vec<_>>();
70assert_eq!(v, it);
71```
72
73or structs:
74
75```rust
76# use rustcomp::rcomp;
77struct Point {
78  x: i32,
79  y: i32,
80}
81#
82# impl Point {
83#    fn new(x: i32, y: i32) -> Self {
84#       Self { x, y }
85#   }
86# }
87
88let points = vec![Point::new(1, 2), Point::new(3, 4), Point::new(5, 6)];
89let v = rcomp![Vec<_>; for Point { x, y } in &points => x + y];
90let it = points.into_iter().map(|Point { x, y }| x + y).collect::<Vec<_>>();
91assert_eq!(v, it);
92```
93
94# Flattening
95
96Flattening nested iterators is supported up to the recursion
97limit by chaining the `for-in` clauses:
98
99```rust
100# use rustcomp::rcomp;
101let matrix = vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]];
102let v = rcomp![Vec<_>; for row in &matrix, col in row => *col * 2, if *col % 2 == 0];
103// nested loops are a much nicer example than iterators here
104let mut it = Vec::new();
105for row in &matrix {
106    for col in row {
107        if *col % 2 == 0 {
108           it.push(*col * 2);
109       }
110    }
111}
112assert_eq!(v, it);
113```
114
115# Advanced Examples
116
117See the [`rcomp!`] macro documentation for some advanced examples,
118like creating a `HashMap` or `HashSet`.
119
120# Note on Iterator Examples
121
122It's important to note that iterator examples used to test the
123comprehensions are _equivalent_ to the comprehensions, but not
124_identical_. The macro expands to nested chains of `flat_map`
125and `filter_map` calls; the examples are written for clarity
126and to show the order of operations in the comprehension. For
127example, the matrix example from earlier expands to:
128
129```rust
130# use rustcomp::rcomp;
131# let matrix = vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]];
132let v = (&matrix)
133    .into_iter()
134    .flat_map(|row| {
135        row.into_iter().filter_map(|col| {
136            if (*col % 2 == 0) && true {
137                Some((*col * 2))
138            } else {
139                None
140            }
141        })
142    })
143    .collect::<Vec<_>>();
144# let mut it = Vec::new();
145# for row in &matrix {
146#     for col in row {
147#         if *col % 2 == 0 {
148#            it.push(*col * 2);
149#        }
150#     }
151# }
152# assert_eq!(v, it);
153```
154
155Notice the use of `into_iter` in the expansion.
156
157# What about `mapcomp`?
158
159I'm aware of the existence of the [`mapcomp`](https://docs.rs/mapcomp/latest/mapcomp/index.html)
160crate, but it differs from this crate in a few ways. For starters,
161`mapcomp` aims to make their syntax as close to Python as possible and
162I think they did a great job; this crate is not trying to do that. The
163goal of this crate is to add comprehensions to Rust in an idiomatic way
164with a syntax that flows naturally with the rest of the language while
165still being concise and powerful. `mapcomp` also provides multiple
166macros for different types of comprehensions while this crate provides
167only one.
168
169On a more technical note, `mapcomp` uses generators internally which was
170okay for Rust 2018, but generators and `yield`-ing are now experimental
171features. This was a big inspiration for this crate, as I wanted to make
172a macro-based solution that didn't require nightly, so I settled on iterators
173in lieu of generators.
174*/
175
176/// Generates an iterator that yields the results of the comprehension. The
177/// syntax allows for flattening, filtering, mapping, and collecting iterators
178/// (in that order).
179///
180/// There are 4 main components to a comprehension:
181/// - The optional collection type, which is passed to a `collect` call by the
182///   macro. If this is omitted, the macro will return an iterator instead of
183///   a collection.
184/// - The `for-in` clause, which iterates over the input(s). This can be
185///   chained (e.g. `for i in v1, j in v2, k in v3, ...`) to flatten nested
186///   iterators, up to the recursion limit.
187/// - The mapping expression, which transforms the input.
188/// - The optional guard expression, which filters the input. Although this is
189///   the last component of the comprehension, it is applied _before_ the
190///   mapping expression. In a sense, the end of the comprehension looks like
191///   a `match` arm. This has a few implications which are explored more in
192///   the [crate-level documentation](crate).
193///
194/// With that explained, here's the full syntax:
195///
196/// ```text
197/// rcomp!([collect_ty;] for <pattern> in <iterator>, ... => <mapper>[, if <guard>]);
198/// ```
199///
200/// # Examples
201///
202/// Comprehensions can be as simple or complex as you want. They can collect
203/// the input, filter it, map it, and flatten it all in one go. For example,
204/// here's how you can create a `HashMap` of numbers and their squares using
205/// a comprehension:
206///
207/// ```rust
208/// # use rustcomp::rcomp;
209/// # use std::collections::HashMap;
210/// let m = rcomp![HashMap<_, _>; for i in 0..10 => (i, i * i)];
211/// # let it = (0..10).map(|i| (i, i * i)).collect::<HashMap<_, _>>();
212/// # assert_eq!(m, it);
213/// ```
214///
215/// Another example is removing duplicates from a `Vec` by converting it to
216/// a `HashSet` and back:
217///
218/// ```rust
219/// # use rustcomp::rcomp;
220/// # use std::collections::HashSet;
221/// let v = vec![1, 2, 3, 4, 5, 1, 2, 3, 4, 5];
222/// let s = rcomp![Vec<_>; for i in rcomp![HashSet<_>; for j in &v => *j] => i];
223/// # let mut s = s;
224/// # let mut it = v
225/// #     .into_iter().collect::<HashSet<_>>()
226/// #     .into_iter().collect::<Vec<_>>();
227/// # it.sort();
228/// # s.sort();
229/// # assert_eq!(s, it);
230/// ```
231///
232/// See the [crate-level documentation](crate) for more examples.
233#[macro_export]
234macro_rules! rcomp {
235    (@__ $($vars:pat),+ in $iter:expr => $mapper:expr $(, if $guard:expr)? $(,)?) => (
236        $iter
237            .into_iter()
238            .filter_map(|$($vars),*| {
239                // `&& true` is a trick to make the guard optional
240                if $($guard &&)? true {
241                    Some($mapper)
242                } else {
243                    None
244                }
245            })
246    );
247    (@__ $($vars:pat),+ in $iter:expr, $($recurse:tt)+) => (
248        $iter
249            .into_iter()
250            .flat_map(|$($vars),*| $crate::rcomp!(@__ $($recurse)+))
251    );
252    // these two rules MUST stay in this order, otherwise the `for`
253    // keyword causes ambiguity. the tt munching shouldn't go too
254    // deep since it has an end condition.
255    (for $($t:tt)*) => (
256        $crate::rcomp!(@__ $($t)*)
257    );
258    ($collect:path; $($t:tt)*) => (
259        $crate::rcomp!($($t)*)
260        .collect::<$collect>()
261    );
262}
263
264#[cfg(test)]
265mod tests {
266    #[test]
267    fn test_vec_comp() {
268        let v: Vec<u32> = vec![1, 2, 3, 4, 5];
269        let expected = v
270            .clone()
271            .into_iter()
272            .filter_map(|i| if i % 2 == 0 { Some(i.into()) } else { None })
273            .collect::<Vec<u64>>();
274        // collect evens
275        let actual = rcomp![Vec<_>; for x in v => u64::from(x), if x % 2 == 0];
276        assert_eq!(expected, actual);
277    }
278
279    #[test]
280    fn test_nested_vec_comp() {
281        let v = vec![vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]];
282        let expected = v
283            .clone()
284            .into_iter()
285            .flatten()
286            .filter(|x| x % 2 == 0)
287            .collect::<Vec<_>>();
288        let actual = rcomp![Vec<_>; for row in v, x in row => x, if x % 2 == 0];
289        assert_eq!(expected, actual);
290    }
291
292    #[test]
293    fn test_full_comprehension() {
294        // essentially a no-op
295        let v: Vec<u32> = vec![1, 2, 3, 4, 5];
296        let expected = v
297            .clone()
298            .into_iter()
299            .filter_map(|i| if i % 2 == 0 { Some(i.into()) } else { None })
300            .collect::<Vec<u64>>();
301        // collect evens
302        let actual = rcomp![Vec<_>; for x in v => u64::from(x), if x % 2 == 0];
303        assert_eq!(expected, actual);
304    }
305
306    #[test]
307    fn test_useless_comp() {
308        let v: Vec<u32> = vec![1, 2, 3, 4, 5];
309        let expected = v.clone().into_iter().collect::<Vec<u32>>();
310        let actual = rcomp![for x in v => x].collect::<Vec<_>>();
311        assert_eq!(expected, actual);
312    }
313
314    #[test]
315    fn test_multiple_vars() {
316        let v: Vec<(i32, i32)> = vec![(1, 2), (3, 4), (5, 6)];
317        let expected: Vec<i32> = v.clone().into_iter().map(|(_, y)| y).collect();
318        let actual = rcomp![for (_, y) in v => y].collect::<Vec<_>>();
319        assert_eq!(expected, actual);
320    }
321}