diff --git a/src/data_traits.rs b/src/data_traits.rs index f43bfb4ef..a8a337fe3 100644 --- a/src/data_traits.rs +++ b/src/data_traits.rs @@ -793,3 +793,77 @@ impl<'a, A: 'a, B: 'a> RawDataSubst for CowRepr<'a, A> } } } + +/// Array representation trait for mapping to an owned array with a different element type. +/// +/// Functions such as [`mapv_into_any`](ArrayBase::mapv_into_any) that alter an array's +/// underlying element type often want to preserve the storage type (e.g., `ArcArray`) +/// of the original array. However, because Rust considers `OwnedRepr` and `OwnedRepr` +/// to be completely different types, a trait must be used to indicate what the mapping is. +/// +/// This trait will map owning storage types to the element-swapped version of themselves; +/// view types are mapped to `OwnedRepr`. The following table summarizes the mappings: +/// +/// | Original Storage Type | Corresponding Array Type | Mapped Storage Type | Output of `from_owned` | +/// | ----------------------- | ------------------------ | ------------------- | ---------------------- | +/// | `OwnedRepr` | `Array` | `OwnedRepr` | `Array` | +/// | `OwnedArcRepr` | `ArcArray` | `OwnedArcRepr` | `ArcArray` | +/// | `CowRepr<'a, A>` | `CowArray<'a, A, D>` | `CowRepr<'a, B>` | `CowArray<'a, B, D>` | +/// | `ViewRepr<&'a mut A>` | `ArrayViewMut<'a, A, D>` | `OwnedRepr` | `Array` | +pub trait DataMappable<'a>: RawData +{ + /// The element-swapped, owning storage representation + type Subst: Data + DataOwned; + + /// Cheaply convert an owned [`Array`] to a different owning array type, as dictated by `Subst`. + /// + /// The owning arrays implement `From`, which is the preferred method for changing the underlying storage. + /// This method (and trait) should be reserved for dealing with changing the element type. + fn from_owned(self_: Array) -> ArrayBase, D>; +} + +impl<'a, A: 'a> DataMappable<'a> for OwnedRepr +{ + type Subst = OwnedRepr; + fn from_owned(self_: Array) -> ArrayBase,D> + { + self_ + } +} + +impl<'a, A: 'a> DataMappable<'a> for OwnedArcRepr +{ + type Subst = OwnedArcRepr; + fn from_owned(self_: Array) -> ArrayBase,D> + { + self_.into() + } +} + + +impl<'a, A: 'a> DataMappable<'a> for CowRepr<'a, A> +{ + type Subst = CowRepr<'a, B>; + fn from_owned(self_: Array) -> ArrayBase,D> + { + self_.into() + } +} + +impl<'a, A: 'a> DataMappable<'a> for ViewRepr<&'a A> +{ + type Subst = OwnedRepr; + fn from_owned(self_: Array) -> ArrayBase, D> + { + self_ + } +} + +impl<'a, A: 'a> DataMappable<'a> for ViewRepr<&'a mut A> +{ + type Subst = OwnedRepr; + fn from_owned(self_: Array) -> ArrayBase, D> + { + self_ + } +} \ No newline at end of file diff --git a/src/impl_methods.rs b/src/impl_methods.rs index 4a00ea000..21342d3ee 100644 --- a/src/impl_methods.rs +++ b/src/impl_methods.rs @@ -17,6 +17,7 @@ use std::mem::{size_of, ManuallyDrop}; use crate::imp_prelude::*; use crate::argument_traits::AssignElem; +use crate::data_traits::DataMappable; use crate::dimension; use crate::dimension::broadcast::co_broadcast; use crate::dimension::reshape_dim; @@ -2803,48 +2804,6 @@ where self } - /// Consume the array, call `f` by **v**alue on each element, and return an - /// owned array with the new values. Works for **any** `F: FnMut(A)->B`. - /// - /// If `A` and `B` are the same type then the map is performed by delegating - /// to [`mapv_into`] and then converting into an owned array. This avoids - /// unnecessary memory allocations in [`mapv`]. - /// - /// If `A` and `B` are different types then a new array is allocated and the - /// map is performed as in [`mapv`]. - /// - /// Elements are visited in arbitrary order. - /// - /// [`mapv_into`]: ArrayBase::mapv_into - /// [`mapv`]: ArrayBase::mapv - pub fn mapv_into_any(self, mut f: F) -> Array - where - S: DataMut, - F: FnMut(A) -> B, - A: Clone + 'static, - B: 'static, - { - if core::any::TypeId::of::() == core::any::TypeId::of::() { - // A and B are the same type. - // Wrap f in a closure of type FnMut(A) -> A . - let f = |a| { - let b = f(a); - // Safe because A and B are the same type. - unsafe { unlimited_transmute::(b) } - }; - // Delegate to mapv_into() using the wrapped closure. - // Convert output to a uniquely owned array of type Array. - let output = self.mapv_into(f).into_owned(); - // Change the return type from Array to Array. - // Again, safe because A and B are the same type. - unsafe { unlimited_transmute::, Array>(output) } - } else { - // A and B are not the same type. - // Fallback to mapv(). - self.mapv(f) - } - } - /// Modify the array in place by calling `f` by mutable reference on each element. /// /// Elements are visited in arbitrary order. @@ -3059,6 +3018,111 @@ where } } +/// # Additional Mapping Methods +impl<'a, A, S, D> ArrayBase +where + D: Dimension, + // Need 'static lifetime bounds for TypeId to work. + // mapv() requires that A be Clone. + A: Clone + 'a + 'static, + // Output is same memory representation as input, substituting B for A. + S: Data + DataMappable<'a>, +{ + /// Consume the array, call `f` by **v**alue on each element, and return an + /// owned array with the new values. Works for **any** `F: FnMut(A)->B`. + /// + /// If `A` and `B` are the same type then the map is performed by delegating + /// to [`mapv_into`](`ArrayBase::mapv_into`) and then converting into an + /// owned array. This avoids unnecessary memory allocations in + /// [`mapv`](`ArrayBase::mapv`). + /// + /// If `A` and `B` are different types then a new array is allocated and the + /// map is performed as in [`mapv`](`ArrayBase::mapv`). + /// + /// Elements are visited in arbitrary order. + /// + /// Example: + /// + /// ```rust + /// # use ndarray::{array, Array}; + /// let a: Array = array![[1., 2., 3.]]; + /// let b = a.clone(); + /// // Same type, no new memory allocation. + /// let a_plus_one = a.mapv_into_any(|a| a + 1.); + /// // Different types, allocates new memory. + /// let rounded = b.mapv_into_any(|a| a.round() as i32); + /// ``` + /// + /// The return data representation/type depends on the input type and is the + /// same as the input type in most cases. See [`DataMappable`] for details. + /// + /// - [`OwnedRepr`](`crate::OwnedRepr`)/[`Array`] -> [`OwnedRepr`](`crate::OwnedRepr`)/[`Array`] + /// - [`OwnedArcRepr`](`crate::OwnedArcRepr`)/[`ArcArray`] -> [`OwnedArcRepr`](`crate::OwnedArcRepr`)/[`ArcArray`] + /// - [`CowRepr`](`crate::CowRepr`)/[`CowArray`] -> [`CowRepr`](`crate::CowRepr`)/[`CowArray`] + /// - [`ViewRepr`](`crate::ViewRepr`)/[`ArrayView`] or [`ArrayViewMut`] -> [`OwnedRepr`](`crate::OwnedRepr`)/[`Array`] + /// + /// Mapping from `A` to a different type `B` will always require new memory + /// to be allocated. Mapping when `A` and `B` are the same type may not need + /// a new memory allocation depending on the input data representation/type. + /// + /// - [`OwnedRepr`](`crate::OwnedRepr`)/[`Array`]: No new memory allocation. + /// - [`OwnedArcRepr`](`crate::OwnedArcRepr`)/[`ArcArray`]: No new memory allocated if data is uniquely owned. + /// - [`CowRepr`](`crate::CowRepr`)/[`CowArray`]: No new memory allocated if data is uniquely owned. + /// - [`ViewRepr`](`crate::ViewRepr`)/[`ArrayView`] or [`ArrayViewMut`]: Always requires new memory allocation. + /// Consider using [`map_inplace`](`ArrayBase::map_inplace`) instead. + /// + /// Example: + /// + /// ```rust + /// # use ndarray::{array, ArcArray}; + /// // Uniquely owned data, no new memory allocation. + /// let a: ArcArray = array![[1., 2., 3.]].into_shared(); + /// let a_plus_one = a.mapv_into_any(|a| a + 1.); + /// // Shared data, requires new memory allocation. + /// let a: ArcArray = array![[1., 2., 3.]].into_shared(); + /// let b = a.clone(); // `a` is shared here + /// let a_plus_one = a.mapv_into_any(|a| a + 1.); // new allocation + /// ``` + /// + /// See also: + /// + /// - [`map_inplace`](`ArrayBase::map_inplace`) + /// - [`mapv_into`](`ArrayBase::mapv_into`) + /// - [`mapv`](`ArrayBase::mapv`) + pub fn mapv_into_any(self, mut f: F) -> ArrayBase<>::Subst, D> + where + // Need 'static lifetime bounds for TypeId to work. + B: 'static, + // Mapping function maps from A to B. + F: FnMut(A) -> B, + { + if core::any::TypeId::of::() == core::any::TypeId::of::() { + // A and B are the same type. + // Wrap f in a closure of type FnMut(A) -> A . + let f = |a| { + let b = f(a); + // Safe because A and B are the same type. + unsafe { unlimited_transmute::(b) } + }; + // Convert to a uniquely-owned data representation: Array + // This will require cloning the data if it is not uniquely owned. + let input = self.into_owned(); + // Delegate to mapv_into() to map from element type A to type A. + let output = input.mapv_into(f); + // Convert to the output data representation, but still with data A. + let output = ::from_owned::(output); + // Transmute to the output representation to data B. + // Safe because A and B are the same type, + // and we are not changing the data representation. + unsafe { unlimited_transmute::::Subst, D>, ArrayBase<::Subst, D>>(output) } + } else { + // A and B are not the same type. + // Fallback to mapv(). + ::from_owned::(self.mapv(f)) + } + } +} + /// Transmute from A to B. /// /// Like transmute, but does not have the compile-time size check which blocks diff --git a/src/lib.rs b/src/lib.rs index f52f25e5e..bf7ee4a32 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -174,7 +174,7 @@ mod data_traits; pub use crate::aliases::*; -pub use crate::data_traits::{Data, DataMut, DataOwned, DataShared, RawData, RawDataClone, RawDataMut, RawDataSubst}; +pub use crate::data_traits::{Data, DataMut, DataOwned, DataShared, RawData, RawDataClone, RawDataMut, RawDataSubst, DataMappable}; mod free_functions; pub use crate::free_functions::*; diff --git a/tests/array.rs b/tests/array.rs index 696904dab..4584f69c0 100644 --- a/tests/array.rs +++ b/tests/array.rs @@ -1053,8 +1053,12 @@ fn map1() fn mapv_into_any_same_type() { let a: Array = array![[1., 2., 3.], [4., 5., 6.]]; + let ptr_before = a.get_ptr((0,0)).unwrap(); let a_plus_one: Array = array![[2., 3., 4.], [5., 6., 7.]]; - assert_eq!(a.mapv_into_any(|a| a + 1.), a_plus_one); + let b = a.mapv_into_any(|a| a + 1.); + let ptr_after = b.get_ptr((0,0)).unwrap(); + assert_eq!(b, a_plus_one); + assert!(ptr_before == ptr_after); // no new memory allocation } #[test] @@ -1062,7 +1066,173 @@ fn mapv_into_any_diff_types() { let a: Array = array![[1., 2., 3.], [4., 5., 6.]]; let a_even: Array = array![[false, true, false], [true, false, true]]; - assert_eq!(a.mapv_into_any(|a| a.round() as i32 % 2 == 0), a_even); + let b = a.mapv_into_any(|a| a.round() as i32 % 2 == 0); + assert_eq!(b, a_even); +} + +#[test] +fn mapv_into_any_arcarray_same_type() { + // Uniquely owned variant. + let a: ArcArray = array![[1., 2., 3.], [4., 5., 6.]].into_shared(); + let ptr_before = a.get_ptr((0,0)).unwrap(); + let a_plus_one: Array = array![[2., 3., 4.], [5., 6., 7.]]; + let b = a.mapv_into_any(|a| a + 1.); + let ptr_after = b.get_ptr((0,0)).unwrap(); + assert_eq!(b, a_plus_one); + assert!(ptr_before == ptr_after); // no new memory allocation + let _ = b.try_into_owned_nocopy().unwrap(); // b should be uniquely owned + // Shared variant. + let aa: ArcArray = array![[1., 2., 3.], [4., 5., 6.]].into_shared(); + let a: ArcArray = aa.clone(); + let ptr_before = a.get_ptr((0,0)).unwrap(); + let b = a.mapv_into_any(|a| a + 1.); + let ptr_after = b.get_ptr((0,0)).unwrap(); + assert_eq!(b, a_plus_one); + assert!(ptr_before != ptr_after); // requires new memory allocation + let _ = b.try_into_owned_nocopy().unwrap(); // b should be uniquely owned +} + +#[test] +fn mapv_into_any_arcarray_diff_types() { + let a: ArcArray = array![[1., 2., 3.], [4., 5., 6.]].into_shared(); + let a_even: Array = array![[false, true, false], [true, false, true]]; + let b = a.mapv_into_any(|a| a.round() as i32 % 2 == 0); + assert_eq!(b, a_even); +} + +#[test] +fn mapv_into_any_cowarray_same_type() { + // Owned variant. + let a: CowArray = array![[1., 2., 3.], [4., 5., 6.]].into(); + assert!(a.is_owned()); + let ptr_before = a.get_ptr((0,0)).unwrap(); + let a_plus_one: Array = array![[2., 3., 4.], [5., 6., 7.]]; + let b = a.mapv_into_any(|a| a + 1.); + let ptr_after = b.get_ptr((0,0)).unwrap(); + assert!(b.is_owned()); + assert_eq!(b, a_plus_one); + assert!(ptr_before == ptr_after); // no new memory allocation + // View variant. + let aa = array![[1., 2., 3.], [4., 5., 6.]]; + let a = CowArray::from(aa.view()); + let ptr_before = a.get_ptr((0,0)).unwrap(); + assert!(a.is_view()); + let b = a.mapv_into_any(|a| a + 1.); + let ptr_after = b.get_ptr((0,0)).unwrap(); + assert!(b.is_owned()); + assert_eq!(b, a_plus_one); + assert!(ptr_before != ptr_after); // requires new memory allocation +} + +#[test] +fn mapv_into_any_cowarray_diff_types() { + // Owned variant. + let a: CowArray = array![[1., 2., 3.], [4., 5., 6.]].into(); + let a_even: Array = array![[false, true, false], [true, false, true]]; + let b = a.mapv_into_any(|a| a.round() as i32 % 2 == 0); + assert_eq!(b, a_even); + // View variant. + let aa = array![[1., 2., 3.], [4., 5., 6.]]; + let a: CowArray = CowArray::from(aa.view()); + let b = a.mapv_into_any(|a| a.round() as i32 % 2 == 0); + assert_eq!(b, a_even); +} + +#[test] +fn mapv_into_any_arrayview_same_type() { + let a: Array = array![[1., 2., 3.], [4., 5., 6.]].into(); + let aview = a.view(); + let a_plus_one: Array = array![[2., 3., 4.], [5., 6., 7.]]; + let b = aview.mapv_into_any(|a| a + 1.); + assert_eq!(b, a_plus_one); +} + +#[test] +fn mapv_into_any_arrayview_diff_types() { + let a: Array = array![[1., 2., 3.], [4., 5., 6.]].into(); + let aview = a.view(); + let a_even: Array = array![[false, true, false], [true, false, true]]; + let b = aview.mapv_into_any(|a| a.round() as i32 % 2 == 0); + assert_eq!(b, a_even); +} + +#[test] +fn mapv_into_any_arrayviewmut_same_type() { + let mut a: Array = array![[1., 2., 3.], [4., 5., 6.]].into(); + let aview = a.view_mut(); + let a_plus_one: Array = array![[2., 3., 4.], [5., 6., 7.]]; + let b = aview.mapv_into_any(|a| a + 1.); + assert_eq!(b, a_plus_one); +} + +#[test] +fn mapv_into_any_arrayviewmut_diff_types() { + let mut a: Array = array![[1., 2., 3.], [4., 5., 6.]].into(); + let aview = a.view_mut(); + let a_even: Array = array![[false, true, false], [true, false, true]]; + let b = aview.mapv_into_any(|a| a.round() as i32 % 2 == 0); + assert_eq!(b, a_even); +} + +use num_traits; +use ndarray::DataMappable; +use ndarray::Data; + +// Test ergonomics on trait bounds of functions calling mapv_into_any(). +// Concrete data representation, generic data type. +fn mapv_into_any_round(arr: Array) -> Array +where + A: num_traits::real::Real + num_traits::cast::AsPrimitive, + D: Dimension, +{ + arr.mapv_into_any(|a| a.round().as_()) +} + +#[test] +fn test_mapv_into_any_round() { + let a: Array = array![1.1, 2.6, 3.0]; + let a_rounded: Array = array![1, 3, 3]; + let b = mapv_into_any_round(a); + assert_eq!(b, a_rounded); +} + +// Test ergonomics on trait bounds of functions calling mapv_into_any(). +// Generic data representation and type. +fn mapv_into_any_round_any<'a, A, S, D>(arr: ArrayBase) -> ArrayBase<>::Subst, D> +where + S: Data + DataMappable<'a>, + A: num_traits::real::Real + num_traits::cast::AsPrimitive + 'a, + D: Dimension, +{ + arr.mapv_into_any(|a| a.round().as_()) +} + +#[test] +fn test_mapv_into_any_round_any() { + let a: Array = array![1.1, 2.6, 3.0]; + let a_rounded: Array = array![1, 3, 3]; + let b = mapv_into_any_round_any(a); + assert_eq!(b, a_rounded); +} + +// Test ergonomics on trait bounds of functions calling mapv_into_any(). +// Generic data representation and type, +// followed by generic reduction. +fn mapv_into_any_round_sum<'a, S, A, D>(arr: ArrayBase) -> i32 +where + S: ndarray::Data + DataMappable<'a>, + A: num_traits::real::Real + num_traits::cast::AsPrimitive + 'a, + D: Dimension, +{ + let output = arr.mapv_into_any(|a| a.round().as_()); + output.sum() +} + +#[test] +fn test_mapv_into_any_round_sum() { + let a: Array = array![1.1, 2.6, 3.0]; + let b = mapv_into_any_round_sum(a); + assert_eq!(b, 7); } #[test]