Merge pull request #797 from rust-ndarray/zip-apply-collect

Add methods to collect Zip into an array

Author: bluss

benches/bench1.rs

@@ -258,23 +258,35 @@ fn add_2d_zip(bench: &mut test::Bencher) {
 }
 
 #[bench]
-fn add_2d_alloc(bench: &mut test::Bencher) {
+fn add_2d_alloc_plus(bench: &mut test::Bencher) {
     let a = Array::<i32, _>::zeros((ADD2DSZ, ADD2DSZ));
     let b = Array::<i32, _>::zeros((ADD2DSZ, ADD2DSZ));
     bench.iter(|| &a + &b);
 }
 
 #[bench]
-fn add_2d_zip_alloc(bench: &mut test::Bencher) {
+fn add_2d_alloc_zip_uninit(bench: &mut test::Bencher) {
     let a = Array::<i32, _>::zeros((ADD2DSZ, ADD2DSZ));
     let b = Array::<i32, _>::zeros((ADD2DSZ, ADD2DSZ));
     bench.iter(|| unsafe {
         let mut c = Array::uninitialized(a.dim());
-        azip!((&a in &a, &b in &b, c in &mut c) *c = a + b);
+        azip!((&a in &a, &b in &b, c in c.raw_view_mut())
+            std::ptr::write(c, a + b)
+        );
         c
     });
 }
 
+#[bench]
+fn add_2d_alloc_zip_collect(bench: &mut test::Bencher) {
+    let a = Array::<i32, _>::zeros((ADD2DSZ, ADD2DSZ));
+    let b = Array::<i32, _>::zeros((ADD2DSZ, ADD2DSZ));
+    bench.iter(|| {
+        Zip::from(&a).and(&b).apply_collect(|&x, &y| x + y)
+    });
+}
+
+
 #[bench]
 fn add_2d_assign_ops(bench: &mut test::Bencher) {
     let mut a = Array::<i32, _>::zeros((ADD2DSZ, ADD2DSZ));

src/argument_traits.rs

@@ -0,0 +1,31 @@
+use std::cell::Cell;
+use std::mem::MaybeUninit;
+
+
+/// A producer element that can be assigned to once
+pub trait AssignElem<T> {
+    /// Assign the value `input` to the element that self represents.
+    fn assign_elem(self, input: T);
+}
+
+/// Assignable element, simply `*self = input`.
+impl<'a, T> AssignElem<T> for &'a mut T {
+    fn assign_elem(self, input: T) {
+        *self = input;
+    }
+}
+
+/// Assignable element, simply `self.set(input)`.
+impl<'a, T> AssignElem<T> for &'a Cell<T> {
+    fn assign_elem(self, input: T) {
+        self.set(input);
+    }
+}
+
+/// Assignable element, the item in the MaybeUninit is overwritten (prior value, if any, is not
+/// read or dropped).
+impl<'a, T> AssignElem<T> for &'a mut MaybeUninit<T> {
+    fn assign_elem(self, input: T) {
+        *self = MaybeUninit::new(input);
+    }
+}

src/impl_constructors.rs

@@ -13,6 +13,7 @@
 #![allow(clippy::match_wild_err_arm)]
 
 use num_traits::{Float, One, Zero};
+use std::mem::MaybeUninit;
 
 use crate::dimension;
 use crate::error::{self, ShapeError};
@@ -517,3 +518,22 @@ where
         Self::from_shape_vec_unchecked(shape, v)
     }
 }
+
+impl<S, A, D> ArrayBase<S, D>
+where
+    S: DataOwned<Elem = MaybeUninit<A>>,
+    D: Dimension,
+{
+    pub(crate) fn maybe_uninit<Sh>(shape: Sh) -> Self
+    where
+        Sh: ShapeBuilder<Dim = D>,
+    {
+        unsafe {
+            let shape = shape.into_shape();
+            let size = size_of_shape_checked_unwrap!(&shape.dim);
+            let mut v = Vec::with_capacity(size);
+            v.set_len(size);
+            Self::from_shape_vec_unchecked(shape, v)
+        }
+    }
+}

src/impl_owned_array.rs

@@ -1,4 +1,8 @@
+use std::mem::MaybeUninit;
+use std::mem::transmute;
+
 use crate::imp_prelude::*;
+use crate::OwnedRepr;
 
 /// Methods specific to `Array0`.
 ///
@@ -57,3 +61,33 @@ where
         self.data.0
     }
 }
+
+/// Methods specific to `Array` of `MaybeUninit`.
+///
+/// ***See also all methods for [`ArrayBase`]***
+///
+/// [`ArrayBase`]: struct.ArrayBase.html
+impl<A, D> Array<MaybeUninit<A>, D>
+where
+    D: Dimension,
+{
+    /// Assert that the array's storage's elements are all fully initialized, and conver
+    /// the array from element type `MaybeUninit<A>` to `A`.
+    pub(crate) unsafe fn assume_init(self) -> Array<A, D> {
+        // NOTE: Fully initialized includes elements not reachable in current slicing/view.
+        //
+        // Should this method be generalized to all array types?
+        // (Will need a way to map the RawData<Elem=X> to RawData<Elem=Y> of same kind)
+
+        let Array { data, ptr, dim, strides } = self;
+        let data = transmute::<OwnedRepr<MaybeUninit<A>>, OwnedRepr<A>>(data);
+        let ptr = ptr.cast::<A>();
+
+        Array {
+            data,
+            ptr,
+            dim,
+            strides,
+        }
+    }
+}

src/layout/layoutfmt.rs

@@ -7,7 +7,6 @@
 // except according to those terms.
 
 use super::Layout;
-use super::LayoutPriv;
 
 const LAYOUT_NAMES: &[&str] = &["C", "F"];
 

src/layout/mod.rs

@@ -1,35 +1,28 @@
 mod layoutfmt;
 
-// public but users don't interact with it
+// public struct but users don't interact with it
 #[doc(hidden)]
 /// Memory layout description
 #[derive(Copy, Clone)]
 pub struct Layout(u32);
 
-pub trait LayoutPriv: Sized {
-    fn new(x: u32) -> Self;
-    fn and(self, flag: Self) -> Self;
-    fn is(self, flag: u32) -> bool;
-    fn flag(self) -> u32;
-}
-
-impl LayoutPriv for Layout {
+impl Layout {
     #[inline(always)]
-    fn new(x: u32) -> Self {
+    pub(crate) fn new(x: u32) -> Self {
         Layout(x)
     }
 
     #[inline(always)]
-    fn is(self, flag: u32) -> bool {
+    pub(crate) fn is(self, flag: u32) -> bool {
         self.0 & flag != 0
     }
     #[inline(always)]
-    fn and(self, flag: Layout) -> Layout {
+    pub(crate) fn and(self, flag: Layout) -> Layout {
         Layout(self.0 & flag.0)
     }
 
     #[inline(always)]
-    fn flag(self) -> u32 {
+    pub(crate) fn flag(self) -> u32 {
         self.0
     }
 }

src/lib.rs

@@ -149,6 +149,8 @@ mod array_approx;
 mod array_serde;
 mod arrayformat;
 mod arraytraits;
+mod argument_traits;
+pub use crate::argument_traits::AssignElem;
 mod data_traits;
 
 pub use crate::aliases::*;

src/parallel/impl_par_methods.rs

@@ -1,4 +1,5 @@
-use crate::{ArrayBase, DataMut, Dimension, NdProducer, Zip};
+use crate::{Array, ArrayBase, DataMut, Dimension, IntoNdProducer, NdProducer, Zip};
+use crate::AssignElem;
 
 use crate::parallel::prelude::*;
 
@@ -43,7 +44,7 @@ where
 // Zip
 
 macro_rules! zip_impl {
-    ($([$($p:ident)*],)+) => {
+    ($([$notlast:ident $($p:ident)*],)+) => {
         $(
         #[allow(non_snake_case)]
         impl<D, $($p),*> Zip<($($p,)*), D>
@@ -63,16 +64,52 @@ macro_rules! zip_impl {
             {
                 self.into_par_iter().for_each(move |($($p,)*)| function($($p),*))
             }
+
+            expand_if!(@bool [$notlast]
+
+            /// Apply and collect the results into a new array, which has the same size as the
+            /// inputs.
+            ///
+            /// If all inputs are c- or f-order respectively, that is preserved in the output.
+            ///
+            /// Restricted to functions that produce copyable results for technical reasons; other
+            /// cases are not yet implemented.
+            pub fn par_apply_collect<R>(self, f: impl Fn($($p::Item,)* ) -> R + Sync + Send) -> Array<R, D>
+                where R: Copy + Send
+            {
+                let mut output = self.uninitalized_for_current_layout::<R>();
+                self.par_apply_assign_into(&mut output, f);
+                unsafe {
+                    output.assume_init()
+                }
+            }
+
+            /// Apply and assign the results into the producer `into`, which should have the same
+            /// size as the other inputs.
+            ///
+            /// The producer should have assignable items as dictated by the `AssignElem` trait,
+            /// for example `&mut R`.
+            pub fn par_apply_assign_into<R, Q>(self, into: Q, f: impl Fn($($p::Item,)* ) -> R + Sync + Send)
+                where Q: IntoNdProducer<Dim=D>,
+                      Q::Item: AssignElem<R> + Send,
+                      Q::Output: Send,
+            {
+                self.and(into)
+                    .par_apply(move |$($p, )* output_| {
+                        output_.assign_elem(f($($p ),*));
+                    });
+            }
+            );
         }
         )+
     }
 }
 
 zip_impl! {
-    [P1],
-    [P1 P2],
-    [P1 P2 P3],
-    [P1 P2 P3 P4],
-    [P1 P2 P3 P4 P5],
-    [P1 P2 P3 P4 P5 P6],
+    [true P1],
+    [true P1 P2],
+    [true P1 P2 P3],
+    [true P1 P2 P3 P4],
+    [true P1 P2 P3 P4 P5],
+    [false P1 P2 P3 P4 P5 P6],
 }

src/zip/mod.rs

@@ -9,13 +9,15 @@
 #[macro_use]
 mod zipmacro;
 
+use std::mem::MaybeUninit;
+
 use crate::imp_prelude::*;
+use crate::AssignElem;
 use crate::IntoDimension;
 use crate::Layout;
 use crate::NdIndex;
 
 use crate::indexes::{indices, Indices};
-use crate::layout::LayoutPriv;
 use crate::layout::{CORDER, FORDER};
 
 /// Return if the expression is a break value.
@@ -579,6 +581,7 @@ pub struct Zip<Parts, D> {
     layout: Layout,
 }
 
+
 impl<P, D> Zip<(P,), D>
 where
     D: Dimension,
@@ -735,6 +738,12 @@ where
         self.dimension[unroll_axis] = inner_len;
         FoldWhile::Continue(acc)
     }
+
+    pub(crate) fn uninitalized_for_current_layout<T>(&self) -> Array<MaybeUninit<T>, D>
+    {
+        let is_f = !self.layout.is(CORDER) && self.layout.is(FORDER);
+        Array::maybe_uninit(self.dimension.clone().set_f(is_f))
+    }
 }
 
 /*
@@ -982,6 +991,42 @@ macro_rules! map_impl {
                     dimension: self.dimension,
                 }
             }
+
+            /// Apply and collect the results into a new array, which has the same size as the
+            /// inputs.
+            ///
+            /// If all inputs are c- or f-order respectively, that is preserved in the output.
+            ///
+            /// Restricted to functions that produce copyable results for technical reasons; other
+            /// cases are not yet implemented.
+            pub fn apply_collect<R>(self, f: impl FnMut($($p::Item,)* ) -> R) -> Array<R, D>
+                where R: Copy,
+            {
+                // To support non-Copy elements, implementation of dropping partial array (on
+                // panic) is needed
+                let mut output = self.uninitalized_for_current_layout::<R>();
+                self.apply_assign_into(&mut output, f);
+                unsafe {
+                    output.assume_init()
+                }
+            }
+
+            /// Apply and assign the results into the producer `into`, which should have the same
+            /// size as the other inputs.
+            ///
+            /// The producer should have assignable items as dictated by the `AssignElem` trait,
+            /// for example `&mut R`.
+            pub fn apply_assign_into<R, Q>(self, into: Q, mut f: impl FnMut($($p::Item,)* ) -> R)
+                where Q: IntoNdProducer<Dim=D>,
+                      Q::Item: AssignElem<R>
+            {
+                self.and(into)
+                    .apply(move |$($p, )* output_| {
+                        output_.assign_elem(f($($p ),*));
+                    });
+            }
+
+
             );
 
             /// Split the `Zip` evenly in two.

src/zip/zipmacro.rs

@@ -122,6 +122,12 @@ macro_rules! azip {
             $(.and($prod))*
             .$apply(|$first_pat, $($pat),*| $body)
     };
+
+    // Unindexed with one or more producer, no loop body
+    (@build $apply:ident $first_prod:expr $(, $prod:expr)* $(,)?) => {
+        $crate::Zip::from($first_prod)
+            $(.and($prod))*
+    };
     // catch-all rule
     (@build $($t:tt)*) => { compile_error!("Invalid syntax in azip!()") };
     ($($t:tt)*) => {