Add try_into_owned_nocopy method

Author: jturner314

src/data_traits.rs

@@ -109,6 +109,15 @@ pub unsafe trait Data: RawData {
         Self::Elem: Clone,
         D: Dimension;
 
+    /// Converts the array into `Array<A, D>` if this is possible without
+    /// cloning the array elements. Otherwise, returns `self_` unchanged.
+    #[doc(hidden)]
+    fn try_into_owned_nocopy<D>(
+        self_: ArrayBase<Self, D>,
+    ) -> Result<ArrayBase<OwnedRepr<Self::Elem>, D>, ArrayBase<Self, D>>
+    where
+        D: Dimension;
+
     /// Return a shared ownership (copy on write) array based on the existing one,
     /// cloning elements if necessary.
     #[doc(hidden)]
@@ -276,6 +285,27 @@ unsafe impl<A> Data for OwnedArcRepr<A> {
         }
     }
 
+    fn try_into_owned_nocopy<D>(
+        self_: ArrayBase<Self, D>,
+    ) -> Result<ArrayBase<OwnedRepr<Self::Elem>, D>, ArrayBase<Self, D>>
+    where
+        D: Dimension,
+    {
+        match Arc::try_unwrap(self_.data.0) {
+            Ok(owned_data) => unsafe {
+                // Safe because the data is equivalent.
+                Ok(ArrayBase::from_data_ptr(owned_data, self_.ptr)
+                    .with_strides_dim(self_.strides, self_.dim))
+            },
+            Err(arc_data) => unsafe {
+                // Safe because the data is equivalent; we're just
+                // reconstructing `self_`.
+                Err(ArrayBase::from_data_ptr(OwnedArcRepr(arc_data), self_.ptr)
+                    .with_strides_dim(self_.strides, self_.dim))
+            },
+        }
+    }
+
     #[allow(clippy::wrong_self_convention)]
     fn to_shared<D>(self_: &ArrayBase<Self, D>) -> ArrayBase<OwnedArcRepr<Self::Elem>, D>
     where
@@ -337,6 +367,16 @@ unsafe impl<A> Data for OwnedRepr<A> {
     {
         self_
     }
+
+    #[inline]
+    fn try_into_owned_nocopy<D>(
+        self_: ArrayBase<Self, D>,
+    ) -> Result<ArrayBase<Self, D>, ArrayBase<Self, D>>
+    where
+        D: Dimension,
+    {
+        Ok(self_)
+    }
 }
 
 unsafe impl<A> DataMut for OwnedRepr<A> {}
@@ -393,6 +433,15 @@ unsafe impl<'a, A> Data for ViewRepr<&'a A> {
     {
         self_.to_owned()
     }
+
+    fn try_into_owned_nocopy<D>(
+        self_: ArrayBase<Self, D>,
+    ) -> Result<ArrayBase<OwnedRepr<Self::Elem>, D>, ArrayBase<Self, D>>
+    where
+        D: Dimension,
+    {
+        Err(self_)
+    }
 }
 
 unsafe impl<'a, A> RawDataClone for ViewRepr<&'a A> {
@@ -438,6 +487,15 @@ unsafe impl<'a, A> Data for ViewRepr<&'a mut A> {
     {
         self_.to_owned()
     }
+
+    fn try_into_owned_nocopy<D>(
+        self_: ArrayBase<Self, D>,
+    ) -> Result<ArrayBase<OwnedRepr<Self::Elem>, D>, ArrayBase<Self, D>>
+    where
+        D: Dimension,
+    {
+        Err(self_)
+    }
 }
 
 unsafe impl<'a, A> DataMut for ViewRepr<&'a mut A> {}
@@ -609,6 +667,22 @@ unsafe impl<'a, A> Data for CowRepr<'a, A> {
             },
         }
     }
+
+    fn try_into_owned_nocopy<D>(
+        self_: ArrayBase<Self, D>,
+    ) -> Result<ArrayBase<OwnedRepr<Self::Elem>, D>, ArrayBase<Self, D>>
+    where
+        D: Dimension,
+    {
+        match self_.data {
+            CowRepr::View(_) => Err(self_),
+            CowRepr::Owned(data) => unsafe {
+                // safe because the data is equivalent so ptr, dims remain valid
+                Ok(ArrayBase::from_data_ptr(data, self_.ptr)
+                    .with_strides_dim(self_.strides, self_.dim))
+            },
+        }
+    }
 }
 
 unsafe impl<'a, A> DataMut for CowRepr<'a, A> where A: Clone {}

src/impl_methods.rs

@@ -241,6 +241,34 @@ where
         S::into_owned(self)
     }
 
+    /// Converts the array into `Array<A, D>` if this is possible without
+    /// cloning the array elements. Otherwise, returns `self` unchanged.
+    ///
+    /// ```
+    /// use ndarray::{array, rcarr2, ArcArray2, Array2};
+    ///
+    /// // Reference-counted, clone-on-write `ArcArray`.
+    /// let a: ArcArray2<_> = rcarr2(&[[1., 2.], [3., 4.]]);
+    /// {
+    ///     // Another reference to the same data.
+    ///     let b: ArcArray2<_> = a.clone();
+    ///     // Since there are two references to the same data, `.into_owned()`
+    ///     // would require cloning the data, so `.try_into_owned_nocopy()`
+    ///     // returns `Err`.
+    ///     assert!(b.try_into_owned_nocopy().is_err());
+    /// }
+    /// // Here, since the second reference has been dropped, the `ArcArray`
+    /// // can be converted into an `Array` without cloning the data.
+    /// let unique: Array2<_> = a.try_into_owned_nocopy().unwrap();
+    /// assert_eq!(unique, array![[1., 2.], [3., 4.]]);
+    /// ```
+    pub fn try_into_owned_nocopy(self) -> Result<Array<A, D>, Self>
+    where
+        S: Data,
+    {
+        S::try_into_owned_nocopy(self)
+    }
+
     /// Turn the array into a shared ownership (copy on write) array,
     /// without any copying.
     pub fn into_shared(self) -> ArcArray<A, D>