Skip to content

Commit

Permalink
Merge pull request #722 from rust-ndarray/update-par-azip
Browse files Browse the repository at this point in the history
Update `par_azip!()` syntax to match `azip!()`
  • Loading branch information
bluss authored Sep 20, 2019
2 parents 0446d5b + 0c1a3d7 commit 5ab292d
Show file tree
Hide file tree
Showing 5 changed files with 40 additions and 76 deletions.
2 changes: 1 addition & 1 deletion benches/par_rayon.rs
Original file line number Diff line number Diff line change
Expand Up @@ -131,7 +131,7 @@ fn rayon_add(bench: &mut Bencher) {
let c = Array2::<f64>::zeros((ADDN, ADDN));
let d = Array2::<f64>::zeros((ADDN, ADDN));
bench.iter(|| {
par_azip!(mut a, b, c, d in {
par_azip!((a in &mut a, b in &b, c in &c, d in &d) {
*a += b.exp() + c.exp() + d.exp();
});
});
Expand Down
57 changes: 6 additions & 51 deletions src/parallel/zipmacro.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,12 @@
/// This is a version of the [`azip`] macro that requires the crate feature
/// `rayon` to be enabled.
///
/// See the [`azip`] macro for more details about the macro syntax!
///
/// This example:
///
/// ```rust,ignore
/// par_azip!(mut a, b, c in { *a = b + c })
/// par_azip!((a in &mut a, &b in &b, &c in &c) { *a = b + c })
/// ```
///
/// Is equivalent to:
Expand Down Expand Up @@ -47,60 +49,13 @@
/// // Compute a simple ternary operation:
/// // elementwise addition of b and c, stored in a
///
/// par_azip!(mut a, b, c in { *a = b + c });
/// par_azip!((a in &mut a, &b in &b, &c in &c) *a = b + c);
///
/// assert_eq!(a, &b + &c);
/// }
/// ```
macro_rules! par_azip {
// Build Zip Rule (index)
(@parse [index => $a:expr, $($aa:expr,)*] $t1:tt in $t2:tt) => {
$crate::par_azip!(@finish ($crate::Zip::indexed($a)) [$($aa,)*] $t1 in $t2)
};
// Build Zip Rule (no index)
(@parse [$a:expr, $($aa:expr,)*] $t1:tt in $t2:tt) => {
$crate::par_azip!(@finish ($crate::Zip::from($a)) [$($aa,)*] $t1 in $t2)
};
// Build Finish Rule (both)
(@finish ($z:expr) [$($aa:expr,)*] [$($p:pat,)+] in { $($t:tt)*}) => {
use $crate::parallel::prelude::*;
#[allow(unused_mut)]
($z)
$(
.and($aa)
)*
.par_apply(|$($p),+| {
$($t)*
})
};
// parsing stack: [expressions] [patterns] (one per operand)
// index uses empty [] -- must be first
(@parse [] [] index $i:pat, $($t:tt)*) => {
$crate::par_azip!(@parse [index =>] [$i,] $($t)*);
};
(@parse [$($exprs:tt)*] [$($pats:tt)*] mut $x:ident ($e:expr) $($t:tt)*) => {
$crate::par_azip!(@parse [$($exprs)* $e,] [$($pats)* mut $x,] $($t)*);
};
(@parse [$($exprs:tt)*] [$($pats:tt)*] mut $x:ident $($t:tt)*) => {
$crate::par_azip!(@parse [$($exprs)* &mut $x,] [$($pats)* mut $x,] $($t)*);
};
(@parse [$($exprs:tt)*] [$($pats:tt)*] , $($t:tt)*) => {
$crate::par_azip!(@parse [$($exprs)*] [$($pats)*] $($t)*);
};
(@parse [$($exprs:tt)*] [$($pats:tt)*] ref $x:ident ($e:expr) $($t:tt)*) => {
$crate::par_azip!(@parse [$($exprs)* $e,] [$($pats)* $x,] $($t)*);
};
(@parse [$($exprs:tt)*] [$($pats:tt)*] ref $x:ident $($t:tt)*) => {
$crate::par_azip!(@parse [$($exprs)* &$x,] [$($pats)* $x,] $($t)*);
};
(@parse [$($exprs:tt)*] [$($pats:tt)*] $x:ident ($e:expr) $($t:tt)*) => {
$crate::par_azip!(@parse [$($exprs)* $e,] [$($pats)* &$x,] $($t)*);
};
(@parse [$($exprs:tt)*] [$($pats:tt)*] $x:ident $($t:tt)*) => {
$crate::par_azip!(@parse [$($exprs)* &$x,] [$($pats)* &$x,] $($t)*);
};
(@parse [$($exprs:tt)*] [$($pats:tt)*] $($t:tt)*) => { };
($($t:tt)*) => {
$crate::par_azip!(@parse [] [] $($t)*);
}
$crate::azip!(@build par_apply $($t)*)
};
}
36 changes: 17 additions & 19 deletions src/zip/zipmacro.rs
Original file line number Diff line number Diff line change
Expand Up @@ -101,32 +101,30 @@
/// ```
#[macro_export]
macro_rules! azip {
// Indexed with a single producer and no trailing comma.
((index $index:pat, $first_pat:pat in $first_prod:expr) $body:expr) => {
$crate::Zip::indexed($first_prod).apply(|$index, $first_pat| $body)
// Indexed with a single producer
// we allow an optional trailing comma after the producers in each rule.
(@build $apply:ident (index $index:pat, $first_pat:pat in $first_prod:expr $(,)?) $body:expr) => {
$crate::Zip::indexed($first_prod).$apply(|$index, $first_pat| $body)
};
// Indexed with more than one producer and no trailing comma.
((index $index:pat, $first_pat:pat in $first_prod:expr, $($pat:pat in $prod:expr),*) $body:expr) => {
// Indexed with more than one producer
(@build $apply:ident (index $index:pat, $first_pat:pat in $first_prod:expr, $($pat:pat in $prod:expr),* $(,)?) $body:expr) => {
$crate::Zip::indexed($first_prod)
$(.and($prod))*
.apply(|$index, $first_pat, $($pat),*| $body)
.$apply(|$index, $first_pat, $($pat),*| $body)
};
// Indexed with trailing comma.
((index $index:pat, $($pat:pat in $prod:expr),+,) $body:expr) => {
azip!((index $index, $($pat in $prod),+) $body)
// Unindexed with a single producer
(@build $apply:ident ($first_pat:pat in $first_prod:expr $(,)?) $body:expr) => {
$crate::Zip::from($first_prod).$apply(|$first_pat| $body)
};
// Unindexed with a single producer and no trailing comma.
(($first_pat:pat in $first_prod:expr) $body:expr) => {
$crate::Zip::from($first_prod).apply(|$first_pat| $body)
};
// Unindexed with more than one producer and no trailing comma.
(($first_pat:pat in $first_prod:expr, $($pat:pat in $prod:expr),*) $body:expr) => {
// Unindexed with more than one producer
(@build $apply:ident ($first_pat:pat in $first_prod:expr, $($pat:pat in $prod:expr),* $(,)?) $body:expr) => {
$crate::Zip::from($first_prod)
$(.and($prod))*
.apply(|$first_pat, $($pat),*| $body)
.$apply(|$first_pat, $($pat),*| $body)
};
// Unindexed with trailing comma.
(($($pat:pat in $prod:expr),+,) $body:expr) => {
azip!(($($pat in $prod),+) $body)
// catch-all rule
(@build $($t:tt)*) => { compile_error!("Invalid syntax in azip!()") };
($($t:tt)*) => {
$crate::azip!(@build apply $($t)*)
};
}
10 changes: 10 additions & 0 deletions tests/azip.rs
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,16 @@ fn test_azip2_3() {
assert!(a != b);
}

#[test]
fn test_azip_syntax_trailing_comma() {
let mut b = Array::<i32, _>::zeros((5, 5));
let mut c = Array::<i32, _>::ones((5, 5));
let a = b.clone();
azip!((b in &mut b, c in &mut c, ) swap(b, c));
assert_eq!(a, c);
assert!(a != b);
}

#[test]
#[cfg(feature = "approx")]
fn test_azip2_sum() {
Expand Down
11 changes: 6 additions & 5 deletions tests/par_azip.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#![cfg(feature = "rayon")]

#[cfg(feature = "approx")]
use itertools::enumerate;
use ndarray::parallel::prelude::*;
use ndarray::prelude::*;
Expand All @@ -9,15 +10,15 @@ use std::sync::atomic::{AtomicUsize, Ordering};
fn test_par_azip1() {
let mut a = Array::zeros(62);
let b = Array::from_elem(62, 42);
par_azip!(mut a in { *a = 42 });
par_azip!((a in &mut a) { *a = 42 });
assert_eq!(a, b);
}

#[test]
fn test_par_azip2() {
let mut a = Array::zeros((5, 7));
let b = Array::from_shape_fn(a.dim(), |(i, j)| 1. / (i + 2 * j) as f32);
par_azip!(mut a, b in { *a = b; });
par_azip!((a in &mut a, &b in &b, ) *a = b );
assert_eq!(a, b);
}

Expand All @@ -33,7 +34,7 @@ fn test_par_azip3() {
*elt = i as f32;
}

par_azip!(mut a (&mut a[..]), b (&b[..]), mut c (&mut c[..]) in {
par_azip!((a in &mut a[..], &b in &b[..], c in &mut c[..]) {
*a += b / 10.;
*c = a.sin();
});
Expand All @@ -48,7 +49,7 @@ fn test_zip_dim_mismatch_1() {
let mut d = a.raw_dim();
d[0] += 1;
let b = Array::from_shape_fn(d, |(i, j)| 1. / (i + 2 * j) as f32);
par_azip!(mut a, b in { *a = b; });
par_azip!((a in &mut a, &b in &b) { *a = b; });
}

#[test]
Expand All @@ -59,7 +60,7 @@ fn test_indices_1() {
}

let count = AtomicUsize::new(0);
par_azip!(index i, elt (&a1) in {
par_azip!((index i, &elt in &a1) {
count.fetch_add(1, Ordering::SeqCst);
assert_eq!(elt, i);
});
Expand Down

0 comments on commit 5ab292d

Please sign in to comment.