Add map_mut and map_axis_mut

src/impl_methods.rs

@@ -892,7 +892,7 @@ impl<A, S, D> ArrayBase<S, D> where S: Data<Elem=A>, D: Dimension
     /// **Panics** if any dimension of `chunk_size` is zero<br>
     /// (**Panics** if `D` is `IxDyn` and `chunk_size` does not match the
     /// number of array axes.)
-    pub fn exact_chunks<E>(&self, chunk_size: E) -> ExactChunks<A, D> 
+    pub fn exact_chunks<E>(&self, chunk_size: E) -> ExactChunks<A, D>
         where E: IntoDimension<Dim=D>,
     {
         exact_chunks_of(self.view(), chunk_size)
@@ -930,7 +930,7 @@ impl<A, S, D> ArrayBase<S, D> where S: Data<Elem=A>, D: Dimension
     ///          [6, 6, 7, 7, 8, 8, 0],
     ///          [6, 6, 7, 7, 8, 8, 0]]));
     /// ```
-    pub fn exact_chunks_mut<E>(&mut self, chunk_size: E) -> ExactChunksMut<A, D> 
+    pub fn exact_chunks_mut<E>(&mut self, chunk_size: E) -> ExactChunksMut<A, D>
         where E: IntoDimension<Dim=D>,
               S: DataMut
     {
@@ -941,13 +941,13 @@ impl<A, S, D> ArrayBase<S, D> where S: Data<Elem=A>, D: Dimension
     ///
     /// The windows are all distinct overlapping views of size `window_size`
     /// that fit into the array's shape.
-    /// 
+    ///
     /// Will yield over no elements if window size is larger
     /// than the actual array size of any dimension.
     ///
     /// The produced element is an `ArrayView<A, D>` with exactly the dimension
     /// `window_size`.
-    /// 
+    ///
     /// **Panics** if any dimension of `window_size` is zero.<br>
     /// (**Panics** if `D` is `IxDyn` and `window_size` does not match the
     /// number of array axes.)
@@ -1694,6 +1694,34 @@ impl<A, S, D> ArrayBase<S, D> where S: Data<Elem=A>, D: Dimension
         }
     }
 
+    /// Call `f` on a mutable reference of each element and create a new array
+    /// with the new values.
+    ///
+    /// Elements are visited in arbitrary order.
+    ///
+    /// Return an array with the same shape as `self`.
+    pub fn map_mut<'a, B, F>(&'a mut self, f: F) -> Array<B, D>
+        where F: FnMut(&mut A) -> B,
+              A: 'a,
+              S: DataMut
+    {
+        let dim = self.dim.clone();
+        let strides = self.strides.clone();
+        if self.is_contiguous() {
+            let slc = self.as_slice_memory_order_mut().unwrap();
+            let v = ::iterators::to_vec_mapped(slc.iter_mut(), f);
+            unsafe {
+                ArrayBase::from_shape_vec_unchecked(
+                    dim.strides(strides), v)
+            }
+        } else {
+            let v = ::iterators::to_vec_mapped(self.iter_mut(), f);
+            unsafe {
+                ArrayBase::from_shape_vec_unchecked(dim, v)
+            }
+        }
+    }
+
     /// Call `f` by **v**alue on each element and create a new array
     /// with the new values.
     ///
@@ -1819,4 +1847,32 @@ impl<A, S, D> ArrayBase<S, D> where S: Data<Elem=A>, D: Dimension
             }
         })
     }
+
+    /// Reduce the values along an axis into just one value, producing a new
+    /// array with one less dimension.
+    /// 1-dimensional lanes are passed as mutable references to the reducer,
+    /// allowing for side-effects.
+    ///
+    /// Elements are visited in arbitrary order.
+    ///
+    /// Return the result as an `Array`.
+    ///
+    /// **Panics** if `axis` is out of bounds.
+    pub fn map_axis_mut<'a, B, F>(&'a mut self, axis: Axis, mut mapping: F)
+        -> Array<B, D::Smaller>
+        where D: RemoveAxis,
+              F: FnMut(ArrayViewMut1<'a, A>) -> B,
+              A: 'a + Clone,
+              S: DataMut,
+    {
+        let view_len = self.len_of(axis);
+        let view_stride = self.strides.axis(axis);
+        // use the 0th subview as a map to each 1d array view extended from
+        // the 0th element.
+        self.subview_mut(axis, 0).map_mut(|first_elt: &mut A| {
+            unsafe {
+                mapping(ArrayViewMut::new_(first_elt, Ix1(view_len), Ix1(view_stride)))
+            }
+        })
+    }
 }

src/iterators/mod.rs

@@ -718,7 +718,7 @@ macro_rules! outer_iter_split_at_impl {
             {
                 assert!(index <= self.iter.len);
                 let right_ptr = if index != self.iter.len {
-                    unsafe { self.iter.offset(index) } 
+                    unsafe { self.iter.offset(index) }
                 }
                 else {
                     self.iter.ptr
@@ -1180,9 +1180,11 @@ use indexes::IndicesIterF;
 
 unsafe impl<F> TrustedIterator for Linspace<F> { }
 unsafe impl<'a, A, D> TrustedIterator for Iter<'a, A, D> { }
+unsafe impl<'a, A, D> TrustedIterator for IterMut<'a, A, D> { }
 unsafe impl<I, F> TrustedIterator for std::iter::Map<I, F>
     where I: TrustedIterator { }
 unsafe impl<'a, A> TrustedIterator for slice::Iter<'a, A> { }
+unsafe impl<'a, A> TrustedIterator for slice::IterMut<'a, A> { }
 unsafe impl TrustedIterator for ::std::ops::Range<usize> { }
 // FIXME: These indices iter are dubious -- size needs to be checked up front.
 unsafe impl<D> TrustedIterator for IndicesIter<D> where D: Dimension { }