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#[test]
36fn test_comprehension() {
37    use std::collections::{HashMap, HashSet};
38    const N: usize = 100;
39    assert_eq!(
40        comprehension!(0..N; @HashSet<_>),
41        (0..N).collect::<HashSet<_>>()
42    );
43    assert_eq!(comprehension!(0..N), (0..N).collect::<Vec<_>>());
44    assert_eq!(
45        comprehension!(0..N; @HashMap<_,_>; i => (i, i + i)),
46        (0..N).map(|i| (i, i + i)).collect::<HashMap<_, _>>()
47    );
48    assert_eq!(
49        comprehension!(0..N; i => i + i),
50        (0..N).map(|i| i + i).collect::<Vec<_>>()
51    );
52    assert_eq!(
53        comprehension!(0..N; &i, i % 2 == 0),
54        (0..N).filter(|&i| i % 2 == 0).collect::<Vec<_>>()
55    );
56    assert_eq!(
57        comprehension!(0..N; i => i + i, i % 2 == 0),
58        (0..N)
59            .filter_map(|i| if i % 2 == 0 { Some(i + i) } else { None })
60            .collect::<Vec<_>>()
61    );
62    assert_eq!(
63        comprehension!(0..N; i => i + i, i % 2 == 0, i % 3 == 0),
64        (0..N)
65            .filter_map(|i| if i % 2 == 0 && i % 3 == 0 {
66                Some(i + i)
67            } else {
68                None
69            })
70            .collect::<Vec<_>>()
71    );
72    assert_eq!(
73        comprehension!(0..N; i => i + i, i % 2 == 0, i % 3 == 0, i % 4 == 0),
74        (0..N)
75            .filter_map(|i| if i % 2 == 0 && i % 3 == 0 && i % 4 == 0 {
76                Some(i + i)
77            } else {
78                None
79            })
80            .collect::<Vec<_>>()
81    );
82    assert_eq!(
83        comprehension!(0..N; @HashMap<_,_>; i => (i / 24, i), i % 2 == 0, i % 3 == 0, i % 4 == 0),
84        (0..N)
85            .filter_map(|i| if i % 2 == 0 && i % 3 == 0 && i % 4 == 0 {
86                Some((i / 24, i))
87            } else {
88                None
89            })
90            .collect::<HashMap<_, _>>()
91    );
92}
competitive/tools/iterable.rs

competitive/tools/
iterable.rs
