FEAT: Add Zip::apply_assign_into

Author: bluss

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/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/zip/mod.rs

@@ -10,6 +10,7 @@
 mod zipmacro;
 
 use crate::imp_prelude::*;
+use crate::AssignElem;
 use crate::IntoDimension;
 use crate::Layout;
 use crate::NdIndex;
@@ -579,6 +580,7 @@ pub struct Zip<Parts, D> {
     layout: Layout,
 }
 
+
 impl<P, D> Zip<(P,), D>
 where
     D: Dimension,
@@ -990,21 +992,36 @@ macro_rules! map_impl {
             ///
             /// Restricted to functions that produce copyable results for technical reasons; other
             /// cases are not yet implemented.
-            pub fn apply_collect<R>(self, mut f: impl FnMut($($p::Item,)* ) -> R) -> Array<R, D>
+            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 is_c = self.layout.is(CORDER);
+                let is_f = !is_c && self.layout.is(FORDER);
+                let mut output = Array::maybe_uninit(self.dimension.clone().set_f(is_f));
+                self.apply_assign_into(&mut output, f);
                 unsafe {
-                    let is_c = self.layout.is(CORDER);
-                    let is_f = !is_c && self.layout.is(FORDER);
-                    let mut output = Array::maybe_uninit(self.dimension.clone().set_f(is_f));
-                    self.and(&mut output)
-                        .apply(move |$($p, )* output_| {
-                            std::ptr::write(output_.as_mut_ptr(), f($($p ),*));
-                        });
                     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)*) => {

tests/azip.rs

@@ -77,6 +77,37 @@ fn test_zip_collect() {
     }
 }
 
+#[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));