Initial implementation of ParallelIterator for AxisIter

bluss · bluss · commit 4fa34aec2b85 · 2016-12-19T13:48:28.000+01:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -39,6 +39,8 @@ optional = true
 blas-sys = { version = "0.6.3", optional = true, default-features = false }
 matrixmultiply = { version = "0.1.11" }
 
+rayon = { version = "0.5.0", optional = true, default-features = false }
+
 [dependencies.serde]
 version = "0.8"
 optional = true
diff --git a/src/iterators/mod.rs b/src/iterators/mod.rs
@@ -21,6 +21,9 @@ use super::{
     Axis,
 };
 
+#[cfg(feature = "rayon")]
+mod par;
+
 /// Base for array iterators
 ///
 /// Iterator element type is `&'a A`.
diff --git a/src/iterators/par.rs b/src/iterators/par.rs
@@ -0,0 +1,78 @@
+
+
+use rayon::par_iter::ParallelIterator;
+use rayon::par_iter::IndexedParallelIterator;
+use rayon::par_iter::ExactParallelIterator;
+use rayon::par_iter::BoundedParallelIterator;
+use rayon::par_iter::internal::{Consumer, UnindexedConsumer};
+use rayon::par_iter::internal::bridge;
+use rayon::par_iter::internal::ProducerCallback;
+use rayon::par_iter::internal::Producer;
+
+use super::AxisIter;
+use imp_prelude::*;
+
+
+
+impl<'a, A, D> ParallelIterator for AxisIter<'a, A, D>
+    where D: Dimension,
+          A: Sync,
+{
+    type Item = <Self as Iterator>::Item;
+    fn drive_unindexed<C>(self, consumer: C) -> C::Result
+        where C: UnindexedConsumer<Self::Item>
+    {
+        bridge(self, consumer)
+    }
+}
+
+impl<'a, A, D> IndexedParallelIterator for AxisIter<'a, A, D>
+    where D: Dimension,
+          A: Sync,
+{
+    fn with_producer<Cb>(self, callback: Cb) -> Cb::Output
+        where Cb: ProducerCallback<Self::Item>
+    {
+        callback.callback(self)
+    }
+}
+
+impl<'a, A, D> ExactParallelIterator for AxisIter<'a, A, D>
+    where D: Dimension,
+          A: Sync,
+{
+    fn len(&mut self) -> usize {
+        self.size_hint().0
+    }
+}
+
+impl<'a, A, D> BoundedParallelIterator for AxisIter<'a, A, D>
+    where D: Dimension,
+          A: Sync,
+{
+    fn upper_bound(&mut self) -> usize {
+        ExactParallelIterator::len(self)
+    }
+
+    fn drive<C>(self, consumer: C) -> C::Result
+        where C: Consumer<Self::Item>
+    {
+        bridge(self, consumer)
+    }
+}
+
+// This is the real magic, I guess
+
+impl<'a, A, D> Producer for AxisIter<'a, A, D>
+    where D: Dimension,
+          A: Sync,
+{
+    fn cost(&mut self, len: usize) -> f64 {
+        // FIXME: No idea about what this is
+        len as f64
+    }
+
+    fn split_at(self, i: usize) -> (Self, Self) {
+        self.split_at(i)
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
@@ -77,6 +77,8 @@ extern crate matrixmultiply;
 extern crate itertools;
 extern crate num_traits as libnum;
 extern crate num_complex;
+#[cfg(feature = "rayon")]
+extern crate rayon;
 
 use std::iter::Zip;
 use std::marker::PhantomData;
diff --git a/tests/rayon.rs b/tests/rayon.rs
@@ -0,0 +1,18 @@
+#![cfg(feature = "rayon")]
+
+extern crate rayon;
+extern crate ndarray;
+
+use ndarray::prelude::*;
+
+use rayon::par_iter::ParallelIterator;
+
+#[test]
+fn test_axis_iter() {
+    let mut a = Array2::<u32>::zeros((10240, 10240));
+    for (i, mut v) in a.axis_iter_mut(Axis(0)).enumerate() {
+        v.fill(i as _);
+    }
+    let s = ParallelIterator::map(a.axis_iter(Axis(0)), |x| x.scalar_sum()).sum();
+    assert_eq!(s, a.scalar_sum());
+}
