Repository: rust-ndarray/ndarray (rust)

Issue:
Is it possible to collect a Zip into Array?
Currently if we need to do elementwise operation on array, we can create a result array, and wrap it together with parameters arrays in `Zip` and use `apply` to assign the result. However, since `Zip` have the shape information, can we add a helper function to allocate a result array with correct shape, so that we don't need to create the result array manually?

Pull request description:
Add methods to collect Zip into an array
Add

- Zip::apply_collect
- Zip::apply_assign_into
- Zip::par_apply_collect
- Zip::par_apply_assign_into

Add internal (for now) constructors for `MaybeUninit<T>`-filled arrays.
Restrict to Copy elements to use collect (because we don't implement dropping a
halfway filled array yet).

This is added to all but the last arity of Zip (because the last item is needed for the result that we collect into.

The assignment functions are generalized so that we can assign to `&mut T`, `&mut MaybeUninit<T>` and `&Cell<T>` when we have elements of `T`. This seems powerful, and could be a concept that's explored further (methods .assign() and .fill() come to mind).
A very minor question is what to name and in which module to put the new trait for this - `AssignElem<T>`.

Fixes #755 

Patch:
diff --git a/benches/bench1.rs b/benches/bench1.rs
index 190ab506..d7fa9f16 100644
--- a/benches/bench1.rs
+++ b/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));
diff --git a/src/argument_traits.rs b/src/argument_traits.rs
new file mode 100644
index 00000000..a93e33a1
--- /dev/null
+++ b/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);
+    }
+}
diff --git a/src/impl_constructors.rs b/src/impl_constructors.rs
index 024a1e3b..54b3a8f1 100644
--- a/src/impl_constructors.rs
+++ b/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)
+        }
+    }
+}
diff --git a/src/impl_owned_array.rs b/src/impl_owned_array.rs
index 96075593..174a6703 100644
--- a/src/impl_owned_array.rs
+++ b/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,
+        }
+    }
+}
diff --git a/src/layout/layoutfmt.rs b/src/layout/layoutfmt.rs
index cb799ca8..d5049512 100644
--- a/src/layout/layoutfmt.rs
+++ b/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"];
 
diff --git a/src/layout/mod.rs b/src/layout/mod.rs
index 741a0e05..57983cc3 100644
--- a/src/layout/mod.rs
+++ b/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
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index f9a208df..c8d750d7 100644
--- a/src/lib.rs
+++ b/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::*;
diff --git a/src/parallel/impl_par_methods.rs b/src/parallel/impl_par_methods.rs
index 88fe769b..867cd07a 100644
--- a/src/parallel/impl_par_methods.rs
+++ b/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],
 }
diff --git a/src/zip/mod.rs b/src/zip/mod.rs
index 33317253..0968dff3 100644
--- a/src/zip/mod.rs
+++ b/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.
diff --git a/src/zip/zipmacro.rs b/src/zip/zipmacro.rs
index ea616a05..3ac73b8f 100644
--- a/src/zip/zipmacro.rs
+++ b/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)*) => {


Test patch:
diff --git a/tests/azip.rs b/tests/azip.rs
index 885db85b..90da4185 100644
--- a/tests/azip.rs
+++ b/tests/azip.rs
@@ -49,6 +49,65 @@ fn test_azip2_3() {
     assert!(a != b);
 }
 
+#[test]
+#[cfg(feature = "approx")]
+fn test_zip_collect() {
+    use approx::assert_abs_diff_eq;
+
+    // test Zip::apply_collect and that it preserves c/f layout.
+
+    let b = Array::from_shape_fn((5, 10), |(i, j)| 1. / (i + 2 * j + 1) as f32);
+    let c = Array::from_shape_fn((5, 10), |(i, j)| f32::exp((i + j) as f32));
+
+    {
+        let a = Zip::from(&b).and(&c).apply_collect(|x, y| x + y);
+
+        assert_abs_diff_eq!(a, &b + &c, epsilon = 1e-6);
+        assert_eq!(a.strides(), b.strides());
+    }
+
+    {
+        let b = b.t();
+        let c = c.t();
+
+        let a = Zip::from(&b).and(&c).apply_collect(|x, y| x + y);
+
+        assert_abs_diff_eq!(a, &b + &c, epsilon = 1e-6);
+        assert_eq!(a.strides(), b.strides());
+    }
+}
+
+#[test]
+#[cfg(feature = "approx")]
+fn test_zip_assign_into() {
+    use approx::assert_abs_diff_eq;
+
+    let mut a = Array::<f32, _>::zeros((5, 10));
+    let b = Array::from_shape_fn((5, 10), |(i, j)| 1. / (i + 2 * j + 1) as f32);
+    let c = Array::from_shape_fn((5, 10), |(i, j)| f32::exp((i + j) as f32));
+
+    Zip::from(&b).and(&c).apply_assign_into(&mut a, |x, y| x + y);
+
+    assert_abs_diff_eq!(a, &b + &c, epsilon = 1e-6);
+}
+
+#[test]
+#[cfg(feature = "approx")]
+fn test_zip_assign_into_cell() {
+    use approx::assert_abs_diff_eq;
+    use std::cell::Cell;
+
+    let a = Array::<Cell<f32>, _>::default((5, 10));
+    let b = Array::from_shape_fn((5, 10), |(i, j)| 1. / (i + 2 * j + 1) as f32);
+    let c = Array::from_shape_fn((5, 10), |(i, j)| f32::exp((i + j) as f32));
+
+    Zip::from(&b).and(&c).apply_assign_into(&a, |x, y| x + y);
+    let a2 = a.mapv(|elt| elt.get());
+
+    assert_abs_diff_eq!(a2, &b + &c, epsilon = 1e-6);
+}
+
+
 #[test]
 fn test_azip_syntax_trailing_comma() {
     let mut b = Array::<i32, _>::zeros((5, 5));
diff --git a/tests/par_zip.rs b/tests/par_zip.rs
index f10f3acd..a24dc8e8 100644
--- a/tests/par_zip.rs
+++ b/tests/par_zip.rs
@@ -70,3 +70,45 @@ fn test_zip_index_4() {
         assert_eq!(*elt, j);
     }
 }
+
+#[test]
+#[cfg(feature = "approx")]
+fn test_zip_collect() {
+    use approx::assert_abs_diff_eq;
+
+    // test Zip::apply_collect and that it preserves c/f layout.
+
+    let b = Array::from_shape_fn((M, N), |(i, j)| 1. / (i + 2 * j + 1) as f32);
+    let c = Array::from_shape_fn((M, N), |(i, j)| f32::ln((1 + i + j) as f32));
+
+    {
+        let a = Zip::from(&b).and(&c).par_apply_collect(|x, y| x + y);
+
+        assert_abs_diff_eq!(a, &b + &c, epsilon = 1e-6);
+        assert_eq!(a.strides(), b.strides());
+    }
+
+    {
+        let b = b.t();
+        let c = c.t();
+
+        let a = Zip::from(&b).and(&c).par_apply_collect(|x, y| x + y);
+
+        assert_abs_diff_eq!(a, &b + &c, epsilon = 1e-6);
+        assert_eq!(a.strides(), b.strides());
+    }
+}
+
+#[test]
+#[cfg(feature = "approx")]
+fn test_zip_assign_into() {
+    use approx::assert_abs_diff_eq;
+
+    let mut a = Array::<f32, _>::zeros((M, N));
+    let b = Array::from_shape_fn((M, N), |(i, j)| 1. / (i + 2 * j + 1) as f32);
+    let c = Array::from_shape_fn((M, N), |(i, j)| f32::ln((1 + i + j) as f32));
+
+    Zip::from(&b).and(&c).par_apply_assign_into(&mut a, |x, y| x + y);
+
+    assert_abs_diff_eq!(a, &b + &c, epsilon = 1e-6);
+}