competitive/tools/
iterable.rs

1#[macro_export]
2macro_rules! comprehension {
3    ($it:expr; @$type:ty) => {
4        $it.collect::<$type>()
5    };
6    ($it:expr) => {
7        comprehension![$it; @Vec<_>]
8    };
9    ($it:expr; @$type:ty; $p:pat => $e:expr) => {
10        comprehension![$it.map(|$p| $e); @$type]
11    };
12    ($it:expr; $p:pat => $($t:tt)*) => {
13        comprehension![$it; @Vec<_>; $p => $($t)*]
14    };
15    ($it:expr; $p:pat, $($t:tt)*) => {
16        comprehension![$it; @Vec<_>; $p, $($t)*]
17    };
18    ($it:expr; @$type:ty; $p:pat => $e:expr) => {
19        comprehension![$it; @$type; $p => $e]
20    };
21    ($it:expr; @$type:ty; $p:pat, $b:expr) => {
22        comprehension![$it.filter(|$p| $b); @$type]
23    };
24    ($it:expr; @$type:ty; $p:pat => $e:expr, $b:expr) => {
25        comprehension![$it.filter_map(|$p| if $b { Some($e) } else { None }); @$type]
26    };
27    ($it:expr; @$type:ty; $p:pat => $e:expr, $b1:expr, $b2:expr) => {
28        comprehension![$it; @$type; $p => $e, $b1 & $b2]
29    };
30    ($it:expr; @$type:ty; $p:pat => $e:expr, $b1:expr, $b2:expr, $($t:tt)*) => {
31        comprehension![$it; @$type; $p => $e, $b1 & $b2, $($t)*]
32    };
33}
34
35#[cfg(test)]
36mod tests {
37    #[test]
38    fn test_comprehension() {
39        use std::collections::{HashMap, HashSet};
40        const N: usize = 100;
41        assert_eq!(
42            comprehension!(0..N; @HashSet<_>),
43            (0..N).collect::<HashSet<_>>()
44        );
45        assert_eq!(comprehension!(0..N), (0..N).collect::<Vec<_>>());
46        assert_eq!(
47            comprehension!(0..N; @HashMap<_,_>; i => (i, i + i)),
48            (0..N).map(|i| (i, i + i)).collect::<HashMap<_, _>>()
49        );
50        assert_eq!(
51            comprehension!(0..N; i => i + i),
52            (0..N).map(|i| i + i).collect::<Vec<_>>()
53        );
54        assert_eq!(
55            comprehension!(0..N; &i, i % 2 == 0),
56            (0..N).filter(|&i| i % 2 == 0).collect::<Vec<_>>()
57        );
58        assert_eq!(
59            comprehension!(0..N; i => i + i, i % 2 == 0),
60            (0..N)
61                .filter_map(|i| if i % 2 == 0 { Some(i + i) } else { None })
62                .collect::<Vec<_>>()
63        );
64        assert_eq!(
65            comprehension!(0..N; i => i + i, i % 2 == 0, i % 3 == 0),
66            (0..N)
67                .filter_map(|i| if i % 2 == 0 && i % 3 == 0 {
68                    Some(i + i)
69                } else {
70                    None
71                })
72                .collect::<Vec<_>>()
73        );
74        assert_eq!(
75            comprehension!(0..N; i => i + i, i % 2 == 0, i % 3 == 0, i % 4 == 0),
76            (0..N)
77                .filter_map(|i| if i % 2 == 0 && i % 3 == 0 && i % 4 == 0 {
78                    Some(i + i)
79                } else {
80                    None
81                })
82                .collect::<Vec<_>>()
83        );
84        assert_eq!(
85            comprehension!(0..N; @HashMap<_,_>; i => (i / 24, i), i % 2 == 0, i % 3 == 0, i % 4 == 0),
86            (0..N)
87                .filter_map(|i| if i % 2 == 0 && i % 3 == 0 && i % 4 == 0 {
88                    Some((i / 24, i))
89                } else {
90                    None
91                })
92                .collect::<HashMap<_, _>>()
93        );
94    }
95}
competitive/tools/iterable.rs

competitive/tools/
iterable.rs
