Repository: rust-ndarray/ndarray (rust)

Issue:
Efficiently map array to a (possibly) different type
[`ArrayBase::mapv_into()`](https://docs.rs/ndarray/0.15.3/ndarray/struct.ArrayBase.html#method.mapv_into) consumes self and reuses its memory allocation to compute a new array with the same shape and type.  Avoiding unnecessary allocations is important when mapping very large arrays.

When writing generic methods, it is often necessary to assume that two arrays may be of different types even though we might expect them to be the same type most of the time.  For example, a method that is generic over complex numbers may be called with all real arguments, all complex arguments, or a mixture of real and complex arguments.  This precludes direct use of `mapv_into()` since we cannot assume the arguments are the same type.

Would it be worth adding something like this [`trans_mapv_into()`](https://gist.github.com/benkay86/b9f3eea271999eccf158bfbbcd30b70e), which maps an array with elements of one type to a (possibly) different type?  If the types are the same it delegates to `mapv_into()`, and if the types are different it falls back to using additional allocations.  Type comparison is performed at runtime pending [specialization](https://github.com/rust-lang/rust/issues/31844).

Pull request:
add mapv_into_any(), resolves #1031
First attempt at pull request for #1031.

Since Rust does not yet have [specialization](https://github.com/rust-lang/rfcs/pull/1210), mapping between generic numerical types in ndarray is challenging when the types *might* be different.  For example, we might want to perform on operation on an array of real numbers to get an array that might be real or might be complex.  We can do this now with `ArrayBase::mapv()`, but if it turns out that both the input and output are the same type (e.g. real numbers) then `mapv()` will perform unnecessary allocation of a new array compared to `mapv_into()`.

This PR proposes a generic method `ArrayBase::mapv_into_any()` whose signature accepts a closure mapping between different types `FnMut(A) -> B` but which uses memory-efficient `mapv_into()` when `A` and `B` are the same type.  The type comparison is performed at runtime.  This is the most memory-efficient way to map between arrays of *possibly* different types using stable Rust.

The use of `unsafe` is discussed [on the Rust language forum](https://users.rust-lang.org/t/runtime-specialization-when-types-are-equal/60628).

Patch:
diff --git a/src/impl_methods.rs b/src/impl_methods.rs
index 51c1403a..1042574c 100644
--- a/src/impl_methods.rs
+++ b/src/impl_methods.rs
@@ -2473,6 +2473,45 @@ 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.
+    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.


Test patch:
diff --git a/tests/array.rs b/tests/array.rs
index 37065ba1..d0fc67de 100644
--- a/tests/array.rs
+++ b/tests/array.rs
@@ -989,6 +989,20 @@ fn map1() {
     assert_eq!(a[(0, 0)], *c[(0, 0)]);
 }
 
+#[test]
+fn mapv_into_any_same_type() {
+    let a: Array<f64, _> = array![[1., 2., 3.], [4., 5., 6.]];
+    let a_plus_one: Array<f64, _> = array![[2., 3., 4.], [5., 6., 7.]];
+    assert_eq!(a.mapv_into_any(|a| a + 1.), a_plus_one);
+}
+
+#[test]
+fn mapv_into_any_diff_types() {
+    let a: Array<f64, _> = array![[1., 2., 3.], [4., 5., 6.]];
+    let a_even: Array<bool, _> = array![[false, true, false], [true, false, true]];
+    assert_eq!(a.mapv_into_any(|a| a.round() as i32 % 2 == 0), a_even);
+}
+
 #[test]
 fn as_slice_memory_order_mut_arcarray() {
     // Test that mutation breaks sharing for `ArcArray`.