Repository: rust-ndarray/ndarray (rust)

Issue:
`.as_slice_memory_order_mut()` is very difficult to use correctly in generic code, and there are a couple of bugs as a result
I just realized that `.as_slice_memory_order_mut()` is very difficult to use correctly in code with generic `S: DataMut` for any purpose other than iterating over the elements in arbitrary order. The issue is that `.as_slice_memory_order_mut()` unshares the data, which may change its layout, so strides obtained *before* calling `.as_slice_memory_order_mut()` may not be correct for accessing elements in the returned slice. As a result, the [implementation of `.map_mut()`](https://github.com/rust-ndarray/ndarray/blob/e7600e82e629a54805a9fc33e50f34151c1faa2e/src/impl_methods.rs#L2325-L2328) and the [implementation of `.zip_mut_with_same_shape()`](https://github.com/rust-ndarray/ndarray/blob/e7600e82e629a54805a9fc33e50f34151c1faa2e/src/impl_methods.rs#L2189-L2198) produce incorrect results in some cases for `ArcArray` inputs. Here's an example:

```rust
use ndarray::prelude::*;
use ndarray::ArcArray2;

fn main() {
    // Fortran-layout `ArcArray`.
    let a: ArcArray2<u8> = {
        let orig = array![[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]];
        let mut orig_f = Array2::zeros(orig.raw_dim().f());
        orig_f.assign(&orig);
        orig_f.into_shared()
    };
    assert_eq!(a, array![[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]);
    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, array![[0, 1], [5, 6]]);
    assert_eq!(b.as_slice_memory_order(), Some(&[0, 5, 1, 6][..]));

    // `.map_mut()` is broken in this case, because the data is contiguous and
    // unsharing the data changes the layout. The current implementation of
    // `.map_mut()` uses the strides of `b` from *before* unsharing to construct
    // the new array, but these are no longer correct *after* the unsharing.
    //
    // As a result, this assertion fails using the current implementation.
    assert_eq!(b.map_mut(|&mut x| x + 10), array![[10, 11], [15, 16]]);
}
```

It's also worth noting that the [current implementation of `try_as_slice_memory_order_mut` is questionable](https://github.com/rust-ndarray/ndarray/blob/e7600e82e629a54805a9fc33e50f34151c1faa2e/src/impl_methods.rs#L1552-L1560), because it checks for contiguity before unsharing with `ensure_unique`, and `DataMut` doesn't guarantee that the array after calling `ensure_unique` will still be contiguous (although this is currently the case for all existing storage types).

I think the best solution is to do the following:

1. Add this constraint to `RawDataMut`: "The implementer must ensure that if the input to `try_ensure_unique` is contiguous, then the output will have the same strides as the input."
2. Change the implementation of `RawDataMut` for `OwnedArcRepr` to uphold this constraint.
3. Document that `.as_slice_memory_order_mut()` unshares the data if necessary, but the strides are preserved.

Pull request description:
Guarantee that `.as_slice_memory_order_mut()` preserves strides 
Fixes #1018.

Patch:
diff --git a/src/data_traits.rs b/src/data_traits.rs
index bd540b1b..858dc7c8 100644
--- a/src/data_traits.rs
+++ b/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;
diff --git a/src/impl_methods.rs b/src/impl_methods.rs
index 28098f68..30776ef5 100644
--- a/src/impl_methods.rs
+++ b/src/impl_methods.rs
@@ -217,7 +217,7 @@ where
                 )
             }
         } else {
-            self.map(|x| x.clone())
+            self.map(A::clone)
         }
     }
 
@@ -1536,6 +1536,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,


Test patch:
diff --git a/tests/array.rs b/tests/array.rs
index edd58adb..37065ba1 100644
--- a/tests/array.rs
+++ b/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];