Skip to main content

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}