Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make mapv_into_any() work for ArcArray, resolves #1280 #1283

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 74 additions & 0 deletions src/data_traits.rs
Original file line number Diff line number Diff line change
Expand Up @@ -793,3 +793,77 @@ impl<'a, A: 'a, B: 'a> RawDataSubst<B> 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<A>` and `OwnedRepr<B>`
/// 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<A>` | `Array<A, D>` | `OwnedRepr<B>` | `Array<B, D>` |
/// | `OwnedArcRepr<A>` | `ArcArray<A, D>` | `OwnedArcRepr<B>` | `ArcArray<B, D>` |
/// | `CowRepr<'a, A>` | `CowArray<'a, A, D>` | `CowRepr<'a, B>` | `CowArray<'a, B, D>` |
/// | `ViewRepr<&'a mut A>` | `ArrayViewMut<'a, A, D>` | `OwnedRepr<B>` | `Array<B, D>` |
pub trait DataMappable<'a>: RawData
{
/// The element-swapped, owning storage representation
type Subst<B: 'a>: Data<Elem = B> + DataOwned;

/// Cheaply convert an owned [`Array<B, D>`] 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<B, D: Dimension>(self_: Array<B, D>) -> ArrayBase<Self::Subst<B>, D>;
}

impl<'a, A: 'a> DataMappable<'a> for OwnedRepr<A>
{
type Subst<B: 'a> = OwnedRepr<B>;
fn from_owned<B: 'a, D: Dimension>(self_: Array<B,D>) -> ArrayBase<Self::Subst<B>,D>
{
self_
}
}

impl<'a, A: 'a> DataMappable<'a> for OwnedArcRepr<A>
{
type Subst<B: 'a> = OwnedArcRepr<B>;
fn from_owned<B: 'a, D: Dimension>(self_: Array<B,D>) -> ArrayBase<Self::Subst<B>,D>
{
self_.into()
}
}


impl<'a, A: 'a> DataMappable<'a> for CowRepr<'a, A>
{
type Subst<B: 'a> = CowRepr<'a, B>;
fn from_owned<B: 'a, D: Dimension>(self_: Array<B,D>) -> ArrayBase<Self::Subst<B>,D>
{
self_.into()
}
}

impl<'a, A: 'a> DataMappable<'a> for ViewRepr<&'a A>
{
type Subst<B: 'a> = OwnedRepr<B>;
fn from_owned<B: 'a, D: Dimension>(self_: Array<B, D>) -> ArrayBase<Self::Subst<B>, D>
{
self_
}
}

impl<'a, A: 'a> DataMappable<'a> for ViewRepr<&'a mut A>
{
type Subst<B: 'a> = OwnedRepr<B>;
fn from_owned<B: 'a, D: Dimension>(self_: Array<B, D>) -> ArrayBase<Self::Subst<B>, D>
{
self_
}
}
148 changes: 106 additions & 42 deletions src/impl_methods.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<B, F>(self, mut f: F) -> Array<B, D>
where
S: DataMut,
F: FnMut(A) -> B,
A: Clone + 'static,
B: 'static,
{
if core::any::TypeId::of::<A>() == core::any::TypeId::of::<B>() {
// 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, A>(b) }
};
// Delegate to mapv_into() using the wrapped closure.
// Convert output to a uniquely owned array of type Array<A, D>.
let output = self.mapv_into(f).into_owned();
// Change the return type from Array<A, D> to Array<B, D>.
// Again, safe because A and B are the same type.
unsafe { unlimited_transmute::<Array<A, D>, Array<B, D>>(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.
Expand Down Expand Up @@ -3059,6 +3018,111 @@ where
}
}

/// # Additional Mapping Methods
impl<'a, A, S, D> ArrayBase<S, D>
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<Elem = A> + 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<f32, _> = 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<f32, _> = 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<f32, _> = 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<B, F>(self, mut f: F) -> ArrayBase<<S as DataMappable<'a>>::Subst<B>, 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::<A>() == core::any::TypeId::of::<B>() {
// 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, A>(b) }
};
// Convert to a uniquely-owned data representation: Array<A, D>
// 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 = <S as DataMappable>::from_owned::<A,D>(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::<ArrayBase<<S as DataMappable>::Subst<A>, D>, ArrayBase<<S as DataMappable>::Subst<B>, D>>(output) }
} else {
// A and B are not the same type.
// Fallback to mapv().
<S as DataMappable>::from_owned::<B,D>(self.mapv(f))
}
}
}

/// Transmute from A to B.
///
/// Like transmute, but does not have the compile-time size check which blocks
Expand Down
2 changes: 1 addition & 1 deletion src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::*;
Expand Down
Loading
Loading