nova_snark/spartan/macros.rs
1/// Macros to give syntactic sugar for zipWith pattern and variants.
2///
3/// ```ignore
4/// use crate::spartan::zip_with;
5/// use itertools::Itertools as _; // we use zip_eq to zip!
6/// let v = vec![0, 1, 2];
7/// let w = vec![2, 3, 4];
8/// let y = vec![4, 5, 6];
9///
10/// // Using the `zip_with!` macro to zip three iterators together and apply a closure
11/// // that sums the elements of each iterator.
12/// let res = zip_with!((v.iter(), w.iter(), y.iter()), |a, b, c| a + b + c)
13/// .collect::<Vec<_>>();
14///
15/// println!("{:?}", res); // Output: [6, 9, 12]
16/// ```
17#[macro_export]
18macro_rules! zip_with {
19 // no iterator projection specified: the macro assumes the arguments *are* iterators
20 // ```ignore
21 // zip_with!((iter1, iter2, iter3), |a, b, c| a + b + c) ->
22 // iter1.zip_eq(iter2.zip_eq(iter3)).map(|(a, (b, c))| a + b + c)
23 // ```
24 //
25 // iterator projection specified: use it on each argument
26 // ```ignore
27 // zip_with!(par_iter, (vec1, vec2, vec3), |a, b, c| a + b + c) ->
28 // vec1.par_iter().zip_eq(vec2.par_iter().zip_eq(vec3.par_iter())).map(|(a, (b, c))| a + b + c)
29 // ````
30 ($($f:ident,)? ($e:expr $(, $rest:expr)*), $($move:ident)? |$($i:ident),+ $(,)?| $($work:tt)*) => {{
31 $crate::zip_with!($($f,)? ($e $(, $rest)*), map, $($move)? |$($i),+| $($work)*)
32 }};
33 // no iterator projection specified: the macro assumes the arguments *are* iterators
34 // optional zipping function specified as well: use it instead of map
35 // ```ignore
36 // zip_with!((iter1, iter2, iter3), for_each, |a, b, c| a + b + c) ->
37 // iter1.zip_eq(iter2.zip_eq(iter3)).for_each(|(a, (b, c))| a + b + c)
38 // ```
39 //
40 //
41 // iterator projection specified: use it on each argument
42 // optional zipping function specified as well: use it instead of map
43 // ```ignore
44 // zip_with!(par_iter, (vec1, vec2, vec3), for_each, |a, b, c| a + b + c) ->
45 // vec1.part_iter().zip_eq(vec2.par_iter().zip_eq(vec3.par_iter())).for_each(|(a, (b, c))| a + b + c)
46 // ```
47 ($($f:ident,)? ($e:expr $(, $rest:expr)*), $worker:ident, $($move:ident,)? |$($i:ident),+ $(,)?| $($work:tt)*) => {{
48 $crate::zip_all!($($f,)? ($e $(, $rest)*))
49 .$worker($($move)? |$crate::nested_idents!($($i),+)| {
50 $($work)*
51 })
52 }};
53}
54
55/// Like `zip_with` but use `for_each` instead of `map`.
56#[macro_export]
57macro_rules! zip_with_for_each {
58 // no iterator projection specified: the macro assumes the arguments *are* iterators
59 // ```ignore
60 // zip_with_for_each!((iter1, iter2, iter3), |a, b, c| a + b + c) ->
61 // iter1.zip_eq(iter2.zip_eq(iter3)).for_each(|(a, (b, c))| a + b + c)
62 // ```
63 //
64 // iterator projection specified: use it on each argument
65 // ```ignore
66 // zip_with_for_each!(par_iter, (vec1, vec2, vec3), |a, b, c| a + b + c) ->
67 // vec1.par_iter().zip_eq(vec2.par_iter().zip_eq(vec3.par_iter())).for_each(|(a, (b, c))| a + b + c)
68 // ````
69 ($($f:ident,)? ($e:expr $(, $rest:expr)*), $($move:ident)? |$($i:ident),+ $(,)?| $($work:tt)*) => {{
70 $crate::zip_with!($($f,)? ($e $(, $rest)*), for_each, $($move)? |$($i),+| $($work)*)
71 }};
72}
73
74// Foldright-like nesting for idents (a, b, c) -> (a, (b, c))
75#[doc(hidden)]
76#[macro_export]
77macro_rules! nested_idents {
78 ($a:ident, $b:ident) => {
79 ($a, $b)
80 };
81 ($first:ident, $($rest:ident),+) => {
82 ($first, $crate::nested_idents!($($rest),+))
83 };
84}
85
86// Fold-right like zipping, with an optional function `f` to apply to each argument
87#[doc(hidden)]
88#[macro_export]
89macro_rules! zip_all {
90 (($e:expr,)) => {
91 $e
92 };
93 ($f:ident, ($e:expr,)) => {
94 $e.$f()
95 };
96 ($f:ident, ($first:expr, $second:expr $(, $rest:expr)*)) => {
97 ($first.$f().zip_eq($crate::zip_all!($f, ($second, $( $rest),*))))
98 };
99 (($first:expr, $second:expr $(, $rest:expr)*)) => {
100 ($first.zip_eq($crate::zip_all!(($second, $( $rest),*))))
101 };
102}