Merge pull request #1019 from jturner314/improve-as_slice_memory_order_mut

Guarantee that `.as_slice_memory_order_mut()` preserves strides

Author: bluss

src/data_traits.rs

@@ -50,8 +50,11 @@ pub unsafe trait RawData: Sized {
 pub unsafe trait RawDataMut: RawData {
     /// If possible, ensures that the array has unique access to its data.
     ///
-    /// If `Self` provides safe mutable access to array elements, then it
-    /// **must** panic or ensure that the data is unique.
+    /// The implementer must ensure that if the input is contiguous, then the
+    /// output has the same strides as input.
+    ///
+    /// Additionally, if `Self` provides safe mutable access to array elements,
+    /// then this method **must** panic or ensure that the data is unique.
     #[doc(hidden)]
     fn try_ensure_unique<D>(_: &mut ArrayBase<Self, D>)
     where
@@ -230,14 +233,9 @@ where
             return;
         }
         if self_.dim.size() <= self_.data.0.len() / 2 {
-            // Create a new vec if the current view is less than half of
-            // backing data.
-            unsafe {
-                *self_ = ArrayBase::from_shape_vec_unchecked(
-                    self_.dim.clone(),
-                    self_.iter().cloned().collect(),
-                );
-            }
+            // Clone only the visible elements if the current view is less than
+            // half of backing data.
+            *self_ = self_.to_owned().into_shared();
             return;
         }
         let rcvec = &mut self_.data.0;

src/impl_methods.rs

@@ -217,7 +217,7 @@ where
                 )
             }
         } else {
-            self.map(|x| x.clone())
+            self.map(A::clone)
         }
     }
 
@@ -1548,6 +1548,10 @@ where
 
     /// Return the array’s data as a slice if it is contiguous,
     /// return `None` otherwise.
+    ///
+    /// In the contiguous case, in order to return a unique reference, this
+    /// method unshares the data if necessary, but it preserves the existing
+    /// strides.
     pub fn as_slice_memory_order_mut(&mut self) -> Option<&mut [A]>
     where
         S: DataMut,

tests/array.rs

@@ -990,8 +990,8 @@ fn map1() {
 }
 
 #[test]
-fn as_slice_memory_order() {
-    // test that mutation breaks sharing
+fn as_slice_memory_order_mut_arcarray() {
+    // Test that mutation breaks sharing for `ArcArray`.
     let a = rcarr2(&[[1., 2.], [3., 4.0f32]]);
     let mut b = a.clone();
     for elt in b.as_slice_memory_order_mut().unwrap() {
@@ -1000,6 +1000,38 @@ fn as_slice_memory_order() {
     assert!(a != b, "{:?} != {:?}", a, b);
 }
 
+#[test]
+fn as_slice_memory_order_mut_cowarray() {
+    // Test that mutation breaks sharing for `CowArray`.
+    let a = arr2(&[[1., 2.], [3., 4.0f32]]);
+    let mut b = CowArray::from(a.view());
+    for elt in b.as_slice_memory_order_mut().unwrap() {
+        *elt = 0.;
+    }
+    assert!(a != b, "{:?} != {:?}", a, b);
+}
+
+#[test]
+fn as_slice_memory_order_mut_contiguous_arcarray() {
+    // Test that unsharing preserves the strides in the contiguous case for `ArcArray`.
+    let a = rcarr2(&[[0, 5], [1, 6], [2, 7], [3, 8], [4, 9]]).reversed_axes();
+    let mut b = a.clone().slice_move(s![.., ..2]);
+    assert_eq!(b.strides(), &[1, 2]);
+    b.as_slice_memory_order_mut().unwrap();
+    assert_eq!(b.strides(), &[1, 2]);
+}
+
+#[test]
+fn as_slice_memory_order_mut_contiguous_cowarray() {
+    // Test that unsharing preserves the strides in the contiguous case for `CowArray`.
+    let a = arr2(&[[0, 5], [1, 6], [2, 7], [3, 8], [4, 9]]).reversed_axes();
+    let mut b = CowArray::from(a.slice(s![.., ..2]));
+    assert!(b.is_view());
+    assert_eq!(b.strides(), &[1, 2]);
+    b.as_slice_memory_order_mut().unwrap();
+    assert_eq!(b.strides(), &[1, 2]);
+}
+
 #[test]
 fn array0_into_scalar() {
     // With this kind of setup, the `Array`'s pointer is not the same as the
@@ -1788,6 +1820,33 @@ fn map_memory_order() {
     assert_eq!(amap.strides(), v.strides());
 }
 
+#[test]
+fn map_mut_with_unsharing() {
+    // Fortran-layout `ArcArray`.
+    let a = rcarr2(&[[0, 5], [1, 6], [2, 7], [3, 8], [4, 9]]).reversed_axes();
+    assert_eq!(a.shape(), &[2, 5]);
+    assert_eq!(a.strides(), &[1, 2]);
+    assert_eq!(
+        a.as_slice_memory_order(),
+        Some(&[0, 5, 1, 6, 2, 7, 3, 8, 4, 9][..])
+    );
+
+    // Shared reference of a portion of `a`.
+    let mut b = a.clone().slice_move(s![.., ..2]);
+    assert_eq!(b.shape(), &[2, 2]);
+    assert_eq!(b.strides(), &[1, 2]);
+    assert_eq!(b.as_slice_memory_order(), Some(&[0, 5, 1, 6][..]));
+    assert_eq!(b, array![[0, 1], [5, 6]]);
+
+    // `.map_mut()` unshares the data. Earlier versions of `ndarray` failed
+    // this assertion. See #1018.
+    assert_eq!(b.map_mut(|&mut x| x + 10), array![[10, 11], [15, 16]]);
+
+    // The strides should be preserved.
+    assert_eq!(b.shape(), &[2, 2]);
+    assert_eq!(b.strides(), &[1, 2]);
+}
+
 #[test]
 fn test_view_from_shape() {
     let s = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11];