diff --git a/src/data_traits.rs b/src/data_traits.rs index bd540b1b8..5cdfa5d7a 100644 --- a/src/data_traits.rs +++ b/src/data_traits.rs @@ -16,7 +16,9 @@ use std::ptr::NonNull; use alloc::sync::Arc; use alloc::vec::Vec; -use crate::{ArrayBase, CowRepr, Dimension, OwnedArcRepr, OwnedRepr, RawViewRepr, ViewRepr}; +use crate::{ + ArcArray, Array, ArrayBase, CowRepr, Dimension, OwnedArcRepr, OwnedRepr, RawViewRepr, ViewRepr, +}; /// Array representation trait. /// @@ -97,16 +99,25 @@ pub unsafe trait Data: RawData { /// Converts the array to a uniquely owned array, cloning elements if necessary. #[doc(hidden)] #[allow(clippy::wrong_self_convention)] - fn into_owned(self_: ArrayBase) -> ArrayBase, D> + fn into_owned(self_: ArrayBase) -> Array where Self::Elem: Clone, D: Dimension; + /// Converts the array into `Array` if this is possible without + /// cloning the array elements. Otherwise, returns `self_` unchanged. + #[doc(hidden)] + fn try_into_owned_nocopy( + self_: ArrayBase, + ) -> Result, ArrayBase> + where + D: Dimension; + /// Return a shared ownership (copy on write) array based on the existing one, /// cloning elements if necessary. #[doc(hidden)] #[allow(clippy::wrong_self_convention)] - fn to_shared(self_: &ArrayBase) -> ArrayBase, D> + fn to_shared(self_: &ArrayBase) -> ArcArray where Self::Elem: Clone, D: Dimension, @@ -259,7 +270,7 @@ where } unsafe impl Data for OwnedArcRepr { - fn into_owned(mut self_: ArrayBase) -> ArrayBase, D> + fn into_owned(mut self_: ArrayBase) -> Array where A: Clone, D: Dimension, @@ -273,8 +284,29 @@ unsafe impl Data for OwnedArcRepr { } } + fn try_into_owned_nocopy( + self_: ArrayBase, + ) -> Result, ArrayBase> + 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(self_: &ArrayBase) -> ArrayBase, D> + fn to_shared(self_: &ArrayBase) -> ArcArray where Self::Elem: Clone, D: Dimension, @@ -327,13 +359,23 @@ unsafe impl RawDataMut for OwnedRepr { unsafe impl Data for OwnedRepr { #[inline] - fn into_owned(self_: ArrayBase) -> ArrayBase, D> + fn into_owned(self_: ArrayBase) -> Array where A: Clone, D: Dimension, { self_ } + + #[inline] + fn try_into_owned_nocopy( + self_: ArrayBase, + ) -> Result, ArrayBase> + where + D: Dimension, + { + Ok(self_) + } } unsafe impl DataMut for OwnedRepr {} @@ -383,13 +425,22 @@ unsafe impl<'a, A> RawData for ViewRepr<&'a A> { } unsafe impl<'a, A> Data for ViewRepr<&'a A> { - fn into_owned(self_: ArrayBase) -> ArrayBase, D> + fn into_owned(self_: ArrayBase) -> Array where Self::Elem: Clone, D: Dimension, { self_.to_owned() } + + fn try_into_owned_nocopy( + self_: ArrayBase, + ) -> Result, ArrayBase> + where + D: Dimension, + { + Err(self_) + } } unsafe impl<'a, A> RawDataClone for ViewRepr<&'a A> { @@ -428,13 +479,22 @@ unsafe impl<'a, A> RawDataMut for ViewRepr<&'a mut A> { } unsafe impl<'a, A> Data for ViewRepr<&'a mut A> { - fn into_owned(self_: ArrayBase) -> ArrayBase, D> + fn into_owned(self_: ArrayBase) -> Array where Self::Elem: Clone, D: Dimension, { self_.to_owned() } + + fn try_into_owned_nocopy( + self_: ArrayBase, + ) -> Result, ArrayBase> + where + D: Dimension, + { + Err(self_) + } } unsafe impl<'a, A> DataMut for ViewRepr<&'a mut A> {} @@ -590,7 +650,7 @@ where unsafe impl<'a, A> Data for CowRepr<'a, A> { #[inline] - fn into_owned(self_: ArrayBase, D>) -> ArrayBase, D> + fn into_owned(self_: ArrayBase, D>) -> Array where A: Clone, D: Dimension, @@ -604,6 +664,22 @@ unsafe impl<'a, A> Data for CowRepr<'a, A> { }, } } + + fn try_into_owned_nocopy( + self_: ArrayBase, + ) -> Result, ArrayBase> + 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 {} diff --git a/src/impl_methods.rs b/src/impl_methods.rs index 28098f68f..f0d57e52e 100644 --- a/src/impl_methods.rs +++ b/src/impl_methods.rs @@ -241,6 +241,34 @@ where S::into_owned(self) } + /// Converts the array into `Array` 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, 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