ndarray/zip/zipmacro.rs
1/// Array zip macro: lock step function application across several arrays and
2/// producers.
3///
4/// This is a shorthand for [`Zip`](struct.Zip.html).
5///
6/// This example:
7///
8/// ```rust,ignore
9/// azip!((a in &mut a, &b in &b, &c in &c) *a = b + c);
10/// ```
11///
12/// Is equivalent to:
13///
14/// ```rust,ignore
15/// Zip::from(&mut a).and(&b).and(&c).for_each(|a, &b, &c| {
16/// *a = b + c
17/// });
18/// ```
19///
20/// The syntax is either
21///
22/// `azip!((` *pat* `in` *expr* `,` *[* *pat* `in` *expr* `,` ... *]* `)` *body_expr* `)`
23///
24/// or, to use `Zip::indexed` instead of `Zip::from`,
25///
26/// `azip!((index` *pat* `,` *pat* `in` *expr* `,` *[* *pat* `in` *expr* `,` ... *]* `)` *body_expr* `)`
27///
28/// The *expr* are expressions whose types must implement `IntoNdProducer`, the
29/// *pat* are the patterns of the parameters to the closure called by
30/// `Zip::for_each`, and *body_expr* is the body of the closure called by
31/// `Zip::for_each`. You can think of each *pat* `in` *expr* as being analogous to
32/// the `pat in expr` of a normal loop `for pat in expr { statements }`: a
33/// pattern, followed by `in`, followed by an expression that implements
34/// `IntoNdProducer` (analogous to `IntoIterator` for a `for` loop).
35///
36/// **Panics** if any of the arrays are not of the same shape.
37///
38/// ## Examples
39///
40/// ```rust
41/// use ndarray::{azip, Array1, Array2, Axis};
42///
43/// type M = Array2<f32>;
44///
45/// fn main() {
46/// // Setup example arrays
47/// let mut a = M::zeros((16, 16));
48/// let mut b = M::zeros(a.dim());
49/// let mut c = M::zeros(a.dim());
50///
51/// // assign values
52/// b.fill(1.);
53/// for ((i, j), elt) in c.indexed_iter_mut() {
54/// *elt = (i + 10 * j) as f32;
55/// }
56///
57/// // Example 1: Compute a simple ternary operation:
58/// // elementwise addition of b and c, stored in a
59/// azip!((a in &mut a, &b in &b, &c in &c) *a = b + c);
60///
61/// assert_eq!(a, &b + &c);
62///
63/// // Example 2: azip!() with index
64/// azip!((index (i, j), &b in &b, &c in &c) {
65/// a[[i, j]] = b - c;
66/// });
67///
68/// assert_eq!(a, &b - &c);
69///
70///
71/// // Example 3: azip!() on references
72/// // See the definition of the function below
73/// borrow_multiply(&mut a, &b, &c);
74///
75/// assert_eq!(a, &b * &c);
76///
77///
78/// // Since this function borrows its inputs, the `IntoNdProducer`
79/// // expressions don't need to explicitly include `&mut` or `&`.
80/// fn borrow_multiply(a: &mut M, b: &M, c: &M) {
81/// azip!((a in a, &b in b, &c in c) *a = b * c);
82/// }
83///
84///
85/// // Example 4: using azip!() without dereference in pattern.
86/// //
87/// // Create a new array `totals` with one entry per row of `a`.
88/// // Use azip to traverse the rows of `a` and assign to the corresponding
89/// // entry in `totals` with the sum across each row.
90/// //
91/// // The row is an array view; it doesn't need to be dereferenced.
92/// let mut totals = Array1::zeros(a.nrows());
93/// azip!((totals in &mut totals, row in a.rows()) *totals = row.sum());
94///
95/// // Check the result against the built in `.sum_axis()` along axis 1.
96/// assert_eq!(totals, a.sum_axis(Axis(1)));
97/// }
98///
99/// ```
100#[macro_export]
101macro_rules! azip {
102 // Indexed with a single producer
103 // we allow an optional trailing comma after the producers in each rule.
104 (@build $apply:ident (index $index:pat, $first_pat:pat in $first_prod:expr $(,)?) $body:expr) => {
105 $crate::Zip::indexed($first_prod).$apply(|$index, $first_pat| $body)
106 };
107 // Indexed with more than one producer
108 (@build $apply:ident (index $index:pat, $first_pat:pat in $first_prod:expr, $($pat:pat in $prod:expr),* $(,)?) $body:expr) => {
109 $crate::Zip::indexed($first_prod)
110 $(.and($prod))*
111 .$apply(|$index, $first_pat, $($pat),*| $body)
112 };
113 // Unindexed with a single producer
114 (@build $apply:ident ($first_pat:pat in $first_prod:expr $(,)?) $body:expr) => {
115 $crate::Zip::from($first_prod).$apply(|$first_pat| $body)
116 };
117 // Unindexed with more than one producer
118 (@build $apply:ident ($first_pat:pat in $first_prod:expr, $($pat:pat in $prod:expr),* $(,)?) $body:expr) => {
119 $crate::Zip::from($first_prod)
120 $(.and($prod))*
121 .$apply(|$first_pat, $($pat),*| $body)
122 };
123
124 // Unindexed with one or more producer, no loop body
125 (@build $apply:ident $first_prod:expr $(, $prod:expr)* $(,)?) => {
126 $crate::Zip::from($first_prod)
127 $(.and($prod))*
128 };
129 // catch-all rule
130 (@build $($t:tt)*) => { compile_error!("Invalid syntax in azip!()") };
131 ($($t:tt)*) => {
132 $crate::azip!(@build for_each $($t)*)
133 };
134}