use uninitialized to avoid traversing or cloning

rust-ndarray · bluss · Mar 12, 2021 · Jan 20, 2021 · Jan 23, 2021 · Jan 23, 2021
commit f7c9da16b483d6b4909f43eb055acf1a896e96ab
diff --git a/src/impl_methods.rs b/src/impl_methods.rs
@@ -2037,6 +2037,46 @@ where
         self.zip_mut_with_by_rows(rhs, f);
     }
 
+    /// Traverse two arrays in unspecified order, in lock step,
+    /// calling the closure `f` on each element pair, and put
+    /// the result into the corresponding element of self.
+    pub fn zip_mut_from_pair<B, C, S1, S2, F>(&mut self, lhs: &ArrayBase<S1, D>, rhs: &ArrayBase<S2, D>, f: F, )
+    where
+        S: DataMut,
+        S1: Data<Elem = B>,
+        S2: Data<Elem = C>,
+        F: Fn(&B, &C) -> A,
+    {
+        debug_assert_eq!(self.shape(), lhs.shape());
+        debug_assert_eq!(self.shape(), rhs.shape());
+
+        if self.dim.strides_equivalent(&self.strides, &lhs.strides)
+            && self.dim.strides_equivalent(&self.strides, &rhs.strides)
+        {
+            if let Some(self_s) = self.as_slice_memory_order_mut() {
+                if let Some(lhs_s) = lhs.as_slice_memory_order() {
+                    if let Some(rhs_s) = rhs.as_slice_memory_order() {
+                        for (s, (l, r)) in
+                        self_s.iter_mut().zip(lhs_s.iter().zip(rhs_s)) {
+                            *s = f(&l, &r);
+                        }
+                        return;
+                    }
+                }
+            }
+        }
+
+        // Otherwise, fall back to the outer iter
+        let n = self.ndim();
+        let dim = self.raw_dim();
+        Zip::from(LanesMut::new(self.view_mut(), Axis(n - 1)))
+            .and(Lanes::new(lhs.broadcast_assume(dim.clone()), Axis(n - 1)))
+            .and(Lanes::new(rhs.broadcast_assume(dim), Axis(n - 1)))
+            .for_each(move |s_row, l_row, r_row| {
+                Zip::from(s_row).and(l_row).and(r_row).for_each(|s, a, b| *s = f(a, b))
+            });
+    }
+
     // zip two arrays where they have different layout or strides
     #[inline(always)]
     fn zip_mut_with_by_rows<B, S2, E, F>(&mut self, rhs: &ArrayBase<S2, E>, mut f: F)

diff --git a/src/impl_ops.rs b/src/impl_ops.rs
@@ -56,20 +56,22 @@ macro_rules! impl_binary_op(
 /// between `self` and `rhs`,
 /// and return the result.
 ///
+/// `self` must be an `Array` or `ArcArray`.
+///
 /// If their shapes disagree, `self` is broadcast to their broadcast shape,
 /// cloning the data if needed.
 ///
 /// **Panics** if broadcasting isn’t possible.
 impl<A, B, S, S2, D, E> $trt<ArrayBase<S2, E>> for ArrayBase<S, D>
 where
-    A: Clone + $trt<B, Output=A>,
+    A: Copy + $trt<B, Output=A>,
     B: Clone,
-    S: Data<Elem=A>,
+    S: DataOwned<Elem=A> + DataMut,
     S2: Data<Elem=B>,
     D: Dimension + BroadcastShape<E>,
     E: Dimension,
 {
-    type Output = Array<A, <D as BroadcastShape<E>>::BroadcastOutput>;
+    type Output = ArrayBase<S, <D as BroadcastShape<E>>::BroadcastOutput>;
     fn $mth(self, rhs: ArrayBase<S2, E>) -> Self::Output
     {
         self.$mth(&rhs)
@@ -79,25 +81,46 @@ where
 /// Perform elementwise
 #[doc=$doc]
 /// between reference `self` and `rhs`,
-/// and return the result as a new `Array`.
+/// and return the result.
+///
+/// `rhs` must be an `Array` or `ArcArray`.
 ///
 /// If their shapes disagree, `self` is broadcast to their broadcast shape,
 /// cloning the data if needed.
 ///
 /// **Panics** if broadcasting isn’t possible.
 impl<'a, A, B, S, S2, D, E> $trt<ArrayBase<S2, E>> for &'a ArrayBase<S, D>
 where
-    A: Clone + $trt<B, Output=A>,
-    B: Clone,
+    A: Clone + $trt<B, Output=B>,
+    B: Copy,
     S: Data<Elem=A>,
-    S2: Data<Elem=B>,
-    D: Dimension + BroadcastShape<E>,
-    E: Dimension,
+    S2: DataOwned<Elem=B> + DataMut,
+    D: Dimension,
+    E: Dimension + BroadcastShape<D>,
 {
-    type Output = Array<A, <D as BroadcastShape<E>>::BroadcastOutput>;
+    type Output = ArrayBase<S2, <E as BroadcastShape<D>>::BroadcastOutput>;
     fn $mth(self, rhs: ArrayBase<S2, E>) -> Self::Output
     {
-        self.$mth(&rhs)
+        let shape = rhs.dim.broadcast_shape(&self.dim).unwrap();
+        if shape.slice() == rhs.dim.slice() {
+            let mut out = rhs.into_dimensionality::<<E as BroadcastShape<D>>::BroadcastOutput>().unwrap();
+            out.zip_mut_with(self, |x, y| {
+                *x = y.clone() $operator x.clone();
+            });
+            out
+        } else {
+            // SAFETY: Overwrite all the elements in the array after
+            // it is created via `zip_mut_from_pair`.
+            let mut out = unsafe {
+                Self::Output::uninitialized(shape.clone().into_pattern())
+            };
+            let lhs = self.broadcast(shape.clone()).unwrap();
+            let rhs = rhs.broadcast(shape).unwrap();
+            out.zip_mut_from_pair(&lhs, &rhs, |x, y| {
+                x.clone() $operator y.clone()
+            });
+            out
+        }
     }
 }
 
@@ -106,32 +129,44 @@ where
 /// between `self` and reference `rhs`,
 /// and return the result.
 ///
+/// `rhs` must be an `Array` or `ArcArray`.
+///
 /// If their shapes disagree, `self` is broadcast to their broadcast shape,
 /// cloning the data if needed.
 ///
 /// **Panics** if broadcasting isn’t possible.
 impl<'a, A, B, S, S2, D, E> $trt<&'a ArrayBase<S2, E>> for ArrayBase<S, D>
 where
-    A: Clone + $trt<B, Output=A>,
+    A: Copy + $trt<B, Output=A>,
     B: Clone,
-    S: Data<Elem=A>,
+    S: DataOwned<Elem=A> + DataMut,
     S2: Data<Elem=B>,
     D: Dimension + BroadcastShape<E>,
     E: Dimension,
 {
-    type Output = Array<A, <D as BroadcastShape<E>>::BroadcastOutput>;
+    type Output = ArrayBase<S, <D as BroadcastShape<E>>::BroadcastOutput>;
     fn $mth(self, rhs: &ArrayBase<S2, E>) -> Self::Output
     {
         let shape = self.dim.broadcast_shape(&rhs.dim).unwrap();
-        let mut self_ = if shape.slice() == self.dim.slice() {
-            self.into_owned().into_dimensionality::<<D as BroadcastShape<E>>::BroadcastOutput>().unwrap()
+        if shape.slice() == self.dim.slice() {
+            let mut out = self.into_dimensionality::<<D as BroadcastShape<E>>::BroadcastOutput>().unwrap();
+            out.zip_mut_with(rhs, |x, y| {
+                *x = x.clone() $operator y.clone();
+            });
+            out
         } else {
-            self.broadcast(shape).unwrap().to_owned()
-        };
-        self_.zip_mut_with(rhs, |x, y| {
-            *x = x.clone() $operator y.clone();
-        });
-        self_
+            // SAFETY: Overwrite all the elements in the array after
+            // it is created via `zip_mut_from_pair`.
+            let mut out = unsafe {
+                Self::Output::uninitialized(shape.clone().into_pattern())
+            };
+            let lhs = self.broadcast(shape.clone()).unwrap();
+            let rhs = rhs.broadcast(shape).unwrap();
+            out.zip_mut_from_pair(&lhs, &rhs, |x, y| {
+                x.clone() $operator y.clone()
+            });
+            out
+        }
     }
 }
 
@@ -140,13 +175,13 @@ where
 /// between references `self` and `rhs`,
 /// and return the result as a new `Array`.
 ///
-/// If their shapes disagree, `self` is broadcast to their broadcast shape,
+/// If their shapes disagree, `self` and `rhs` is broadcast to their broadcast shape,
 /// cloning the data if needed.
 ///
 /// **Panics** if broadcasting isn’t possible.
 impl<'a, A, B, S, S2, D, E> $trt<&'a ArrayBase<S2, E>> for &'a ArrayBase<S, D>
 where
-    A: Clone + $trt<B, Output=A>,
+    A: Copy + $trt<B, Output=A>,
     B: Clone,
     S: Data<Elem=A>,
     S2: Data<Elem=B>,
@@ -156,15 +191,17 @@ where
     type Output = Array<A, <D as BroadcastShape<E>>::BroadcastOutput>;
     fn $mth(self, rhs: &'a ArrayBase<S2, E>) -> Self::Output {
         let shape = self.dim.broadcast_shape(&rhs.dim).unwrap();
-        let mut self_ = if shape.slice() == self.dim.slice() {
-            self.to_owned().into_dimensionality::<<D as BroadcastShape<E>>::BroadcastOutput>().unwrap()
-        } else {
-            self.broadcast(shape).unwrap().to_owned()
+        // SAFETY: Overwrite all the elements in the array after
+        // it is created via `zip_mut_from_pair`.
+        let mut out = unsafe {
+            Self::Output::uninitialized(shape.clone().into_pattern())
         };
-        self_.zip_mut_with(rhs, |x, y| {
-            *x = x.clone() $operator y.clone();
+        let lhs = self.broadcast(shape.clone()).unwrap();
+        let rhs = rhs.broadcast(shape).unwrap();
+        out.zip_mut_from_pair(&lhs, &rhs, |x, y| {
+                x.clone() $operator y.clone()
         });
-        self_
+        out
     }
 }
 

diff --git a/src/numeric/impl_numeric.rs b/src/numeric/impl_numeric.rs
@@ -243,7 +243,7 @@ where
     /// **Panics** if `axis` is out of bounds.
     pub fn sum_axis(&self, axis: Axis) -> Array<A, D::Smaller>
     where
-        A: Clone + Zero + Add<Output = A>,
+        A: Copy + Zero + Add<Output = A>,
         D: RemoveAxis,
     {
         let n = self.len_of(axis);
@@ -285,7 +285,7 @@ where
     /// ```
     pub fn mean_axis(&self, axis: Axis) -> Option<Array<A, D::Smaller>>
     where
-        A: Clone + Zero + FromPrimitive + Add<Output = A> + Div<Output = A>,
+        A: Copy + Zero + FromPrimitive + Add<Output = A> + Div<Output = A>,
         D: RemoveAxis,
     {
         let axis_length = self.len_of(axis);

diff --git a/tests/array.rs b/tests/array.rs
@@ -484,7 +484,7 @@ fn test_add() {
     }
 
     let B = A.clone();
-    let A = A + &B;
+    A = A + &B;
     assert_eq!(A[[0, 0]], 0);
     assert_eq!(A[[0, 1]], 2);
     assert_eq!(A[[1, 0]], 4);
@@ -1640,7 +1640,7 @@ fn arithmetic_broadcast() {
         arr3(&[[[11, 15], [20, 24]], [[10, 14], [19, 23]]])
     );
     assert_eq!(
-        &a + b + c.into_owned(),
+        &a + b.into_owned() + c,
         arr3(&[[[15, 19], [32, 36]], [[14, 18], [31, 35]]])
     );
 
@@ -1652,7 +1652,7 @@ fn arithmetic_broadcast() {
     let sc = c.to_shared();
     let sc2 = sc.into_shared();
     assert_eq!(
-        sa2 + sb2 + sc2.into_owned(),
+        sa2 + &sb2 + sc2.into_owned(),
         arr3(&[[[15, 19], [32, 36]], [[14, 18], [31, 35]]])
     );
 }