Merge pull request #947 from rust-ndarray/assign-into

Add new method .assign_to()

Author: bluss

src/impl_constructors.rs

@@ -531,30 +531,14 @@ where
     ///     // two first columns in b are two last in a
     ///     // rest of columns in b are the initial columns in a
     ///
-    ///     assign_to(a.slice(s![.., -2..]), b.slice_mut(s![.., ..2]));
-    ///     assign_to(a.slice(s![.., 2..]), b.slice_mut(s![.., ..-2]));
+    ///     a.slice(s![.., -2..]).assign_to(b.slice_mut(s![.., ..2]));
+    ///     a.slice(s![.., 2..]).assign_to(b.slice_mut(s![.., ..-2]));
     ///
     ///     // Now we can promise that `b` is safe to use with all operations
     ///     unsafe {
     ///         b.assume_init()
     ///     }
     /// }
-    ///
-    /// use ndarray::{IntoNdProducer, AssignElem};
-    ///
-    /// // This function clones elements from the first input to the second;
-    /// // the two producers must have the same shape
-    /// fn assign_to<'a, P1, P2, A>(from: P1, to: P2)
-    ///     where P1: IntoNdProducer<Item = &'a A>,
-    ///           P2: IntoNdProducer<Dim = P1::Dim>,
-    ///           P2::Item: AssignElem<A>,
-    ///           A: Clone + 'a
-    /// {
-    ///     Zip::from(from)
-    ///         .map_assign_into(to, A::clone);
-    /// }
-    ///
-    /// # shift_by_two(&Array2::zeros((8, 8)));
     /// ```
     pub fn uninit<Sh>(shape: Sh) -> ArrayBase<S::MaybeUninit, D>
     where

src/impl_methods.rs

@@ -15,6 +15,7 @@ use rawpointer::PointerExt;
 use crate::imp_prelude::*;
 
 use crate::{arraytraits, DimMax};
+use crate::argument_traits::AssignElem;
 use crate::dimension;
 use crate::dimension::IntoDimension;
 use crate::dimension::{
@@ -25,7 +26,7 @@ use crate::dimension::broadcast::co_broadcast;
 use crate::error::{self, ErrorKind, ShapeError, from_kind};
 use crate::math_cell::MathCell;
 use crate::itertools::zip;
-use crate::zip::Zip;
+use crate::zip::{IntoNdProducer, Zip};
 use crate::AxisDescription;
 
 use crate::iter::{
@@ -2049,6 +2050,23 @@ where
         self.zip_mut_with(rhs, |x, y| *x = y.clone());
     }
 
+    /// Perform an elementwise assigment of values cloned from `self` into array or producer `to`.
+    ///
+    /// The destination `to` can be another array or a producer of assignable elements.
+    /// [`AssignElem`] determines how elements are assigned.
+    ///
+    /// **Panics** if shapes disagree.
+    pub fn assign_to<P>(&self, to: P)
+    where
+        S: Data,
+        P: IntoNdProducer<Dim = D>,
+        P::Item: AssignElem<A>,
+        A: Clone,
+    {
+        Zip::from(self)
+            .map_assign_into(to, A::clone);
+    }
+
     /// Perform an elementwise assigment to `self` from element `x`.
     pub fn fill(&mut self, x: A)
     where

src/lib.rs

@@ -208,7 +208,6 @@ mod shape_builder;
 mod slice;
 mod split_at;
 mod stacking;
-mod traversal_utils;
 #[macro_use]
 mod zip;
 

src/stacking.rs

@@ -8,7 +8,6 @@
 
 use crate::error::{from_kind, ErrorKind, ShapeError};
 use crate::imp_prelude::*;
-use crate::traversal_utils::assign_to;
 
 /// Stack arrays along the new axis.
 ///
@@ -99,7 +98,7 @@ where
         for array in arrays {
             let len = array.len_of(axis);
             let (front, rest) = assign_view.split_at(axis, len);
-            assign_to(array, front);
+            array.assign_to(front);
             assign_view = rest;
         }
         debug_assert_eq!(assign_view.len(), 0);
@@ -171,7 +170,7 @@ where
             // but same number of axes).
             let assign_view = assign_view.into_dimensionality::<D>()
                 .expect("same-dimensionality cast");
-            assign_to(array, assign_view);
+            array.assign_to(assign_view);
         });
 
     unsafe {

src/traversal_utils.rs

@@ -900,31 +900,6 @@ fn standard_layout() {
     assert!(x4.is_standard_layout());
 }
 
-#[test]
-fn assign() {
-    let mut a = arr2(&[[1., 2.], [3., 4.]]);
-    let b = arr2(&[[1., 3.], [2., 4.]]);
-    a.assign(&b);
-    assert_eq!(a, b);
-
-    /* Test broadcasting */
-    a.assign(&ArcArray::zeros(1));
-    assert_eq!(a, ArcArray::zeros((2, 2)));
-
-    /* Test other type */
-    a.assign(&Array::from_elem((2, 2), 3.));
-    assert_eq!(a, ArcArray::from_elem((2, 2), 3.));
-
-    /* Test mut view */
-    let mut a = arr2(&[[1, 2], [3, 4]]);
-    {
-        let mut v = a.view_mut();
-        v.slice_collapse(s![..1, ..]);
-        v.fill(0);
-    }
-    assert_eq!(a, arr2(&[[0, 0], [3, 4]]));
-}
-
 #[test]
 fn iter_size_hint() {
     let mut a = arr2(&[[1., 2.], [3., 4.]]);

tests/array.rs

@@ -6,3 +6,37 @@ fn test_fold_axis_oob() {
     let a = arr2(&[[1., 2.], [3., 4.]]);
     a.fold_axis(Axis(2), 0., |x, y| x + y);
 }
+
+#[test]
+fn assign() {
+    let mut a = arr2(&[[1., 2.], [3., 4.]]);
+    let b = arr2(&[[1., 3.], [2., 4.]]);
+    a.assign(&b);
+    assert_eq!(a, b);
+
+    /* Test broadcasting */
+    a.assign(&ArcArray::zeros(1));
+    assert_eq!(a, ArcArray::zeros((2, 2)));
+
+    /* Test other type */
+    a.assign(&Array::from_elem((2, 2), 3.));
+    assert_eq!(a, ArcArray::from_elem((2, 2), 3.));
+
+    /* Test mut view */
+    let mut a = arr2(&[[1, 2], [3, 4]]);
+    {
+        let mut v = a.view_mut();
+        v.slice_collapse(s![..1, ..]);
+        v.fill(0);
+    }
+    assert_eq!(a, arr2(&[[0, 0], [3, 4]]));
+}
+
+
+#[test]
+fn assign_to() {
+    let mut a = arr2(&[[1., 2.], [3., 4.]]);
+    let b = arr2(&[[0., 3.], [2., 0.]]);
+    b.assign_to(&mut a);
+    assert_eq!(a, b);
+}