Merge pull request #292 from bluss/zip-rows

Reimplement rows case in `zip_mut_with` using `Zip`

Author: bluss

src/impl_methods.rs

@@ -23,11 +23,12 @@ use super::ZipExt;
 use dimension::IntoDimension;
 use dimension::{axes_of, Axes, merge_axes, stride_offset};
 use iterators::{
-    new_inner_iter_smaller,
-    new_inner_iter_smaller_mut,
+    new_inners,
+    new_inners_mut,
     whole_chunks_of,
     whole_chunks_mut_of,
 };
+use zip::Zip;
 
 use {
     NdIndex,
@@ -1184,33 +1185,18 @@ impl<A, S, D> ArrayBase<S, D> where S: Data<Elem=A>, D: Dimension
             }
             return;
         }
+
         // otherwise, break the arrays up into their inner rows
-        let mut try_slices = true;
-        let rows = new_inner_iter_smaller_mut(self.view_mut()).zip(
-                    new_inner_iter_smaller(rhs.view()));
-        for (mut s_row, r_row) in rows {
-            if try_slices {
-                if let Some(self_s) = s_row.as_slice_mut() {
-                    if let Some(rhs_s) = r_row.as_slice() {
-                        let len = cmp::min(self_s.len(), rhs_s.len());
-                        let s = &mut self_s[..len];
-                        let r = &rhs_s[..len];
-                        for i in 0..len {
-                            f(&mut s[i], &r[i]);
-                        }
-                        continue;
-                    }
-                }
-                try_slices = false;
-            }
-            unsafe {
-                for i in 0..s_row.len() {
-                    f(s_row.uget_mut(i), r_row.uget(i))
-                }
-            }
-        }
+        let n = self.ndim();
+        let dim = self.raw_dim();
+        Zip::from(new_inners_mut(self.view_mut(), Axis(n - 1)))
+            .and(new_inners(rhs.broadcast_assume(dim), Axis(n - 1)))
+            .apply(move |s_row, r_row| {
+                Zip::from(s_row).and(r_row).apply(|a, b| f(a, b))
+            });
     }
 
+
     fn zip_mut_with_elem<B, F>(&mut self, rhs_elem: &B, mut f: F)
         where S: DataMut,
               F: FnMut(&mut A, &B)
@@ -1432,39 +1418,4 @@ impl<A, S, D> ArrayBase<S, D> where S: Data<Elem=A>, D: Dimension
             }
         })
     }
-
-    #[cfg(lanes_along)]
-    fn lanes_along<'a, F>(&'a self, axis: Axis, mut visit: F)
-        where D: RemoveAxis,
-              F: FnMut(ArrayView1<'a, A>),
-              A: 'a,
-    {
-        let view_len = self.shape().axis(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(axis, 0).visit(move |first_elt| {
-            unsafe {
-                visit(ArrayView::new_(first_elt, Ix1(view_len), Ix1(view_stride)))
-            }
-        })
-    }
-
-    #[cfg(lanes_along)]
-    fn lanes_along_mut<'a, F>(&'a mut self, axis: Axis, mut visit: F)
-        where D: RemoveAxis,
-              S: DataMut,
-              F: FnMut(ArrayViewMut1<'a, A>),
-              A: 'a,
-    {
-        let view_len = self.shape().axis(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).unordered_foreach_mut(move |first_elt| {
-            unsafe {
-                visit(ArrayViewMut::new_(first_elt, Ix1(view_len), Ix1(view_stride)))
-            }
-        })
-    }
 }

src/iterators/inners.rs

@@ -0,0 +1,101 @@
+
+use imp_prelude::*;
+use {NdProducer, Layout};
+
+impl_ndproducer! {
+    ['a, A, D: Dimension]
+    [Clone => 'a, A, D: Clone ]
+    Inners {
+        base,
+        inner_len,
+        inner_stride,
+    }
+    Inners<'a, A, D> {
+        type Dim = D;
+        type Item = ArrayView<'a, A, Ix1>;
+
+        unsafe fn item(&self, ptr) {
+            ArrayView::new_(ptr, Ix1(self.inner_len), Ix1(self.inner_stride as Ix))
+        }
+    }
+}
+
+pub struct Inners<'a, A: 'a, D> {
+    base: ArrayView<'a, A, D>,
+    inner_len: Ix,
+    inner_stride: Ixs,
+}
+
+
+pub fn new_inners<A, D>(v: ArrayView<A, D>, axis: Axis)
+    -> Inners<A, D::TrySmaller>
+    where D: Dimension
+{
+    let ndim = v.ndim();
+    let len;
+    let stride;
+    let iter_v;
+    if ndim == 0 {
+        len = 1;
+        stride = 0;
+        iter_v = v.try_remove_axis(Axis(0))
+    } else {
+        len = v.dim.last_elem();
+        stride = v.strides.last_elem() as isize;
+        iter_v = v.try_remove_axis(axis)
+    }
+    Inners {
+        inner_len: len,
+        inner_stride: stride,
+        base: iter_v,
+    }
+}
+
+impl_ndproducer! {
+    ['a, A, D: Dimension]
+    [Clone =>]
+    InnersMut {
+        base,
+        inner_len,
+        inner_stride,
+    }
+    InnersMut<'a, A, D> {
+        type Dim = D;
+        type Item = ArrayViewMut<'a, A, Ix1>;
+
+        unsafe fn item(&self, ptr) {
+            ArrayViewMut::new_(ptr, Ix1(self.inner_len), Ix1(self.inner_stride as Ix))
+        }
+    }
+}
+
+pub struct InnersMut<'a, A: 'a, D> {
+    base: ArrayViewMut<'a, A, D>,
+    inner_len: Ix,
+    inner_stride: Ixs,
+}
+
+
+pub fn new_inners_mut<A, D>(v: ArrayViewMut<A, D>, axis: Axis)
+    -> InnersMut<A, D::TrySmaller>
+    where D: Dimension
+{
+    let ndim = v.ndim();
+    let len;
+    let stride;
+    let iter_v;
+    if ndim == 0 {
+        len = 1;
+        stride = 0;
+        iter_v = v.try_remove_axis(Axis(0))
+    } else {
+        len = v.dim.last_elem();
+        stride = v.strides.last_elem() as isize;
+        iter_v = v.try_remove_axis(axis)
+    }
+    InnersMut {
+        inner_len: len,
+        inner_stride: stride,
+        base: iter_v,
+    }
+}

src/iterators/mod.rs

@@ -9,6 +9,7 @@
 
 #[macro_use] mod macros;
 mod chunks;
+mod inners;
 
 use std::marker::PhantomData;
 use std::ptr;
@@ -35,6 +36,10 @@ pub use self::chunks::{
     WholeChunksIterMut,
     whole_chunks_mut_of,
 };
+pub use self::inners::{
+    new_inners,
+    new_inners_mut,
+};
 
 /// Base for array iterators
 ///
@@ -527,6 +532,7 @@ impl<'a, A, D> ExactSizeIterator for InnerIterMut<'a, A, D>
     }
 }
 
+#[cfg(next_version)]
 /// Create an InnerIter one dimension smaller than D (if possible)
 pub fn new_inner_iter_smaller<A, D>(v: ArrayView<A, D>)
     -> InnerIter<A, D::TrySmaller>
@@ -552,6 +558,7 @@ pub fn new_inner_iter_smaller<A, D>(v: ArrayView<A, D>)
     }
 }
 
+#[cfg(next_version)]
 pub fn new_inner_iter_smaller_mut<A, D>(v: ArrayViewMut<A, D>)
     -> InnerIterMut<A, D::TrySmaller>
     where D: Dimension,

src/lib.rs

@@ -590,6 +590,22 @@ impl<A, S, D> ArrayBase<S, D>
         }
     }
 
+    // Broadcast to dimension `E`, without checking that the dimensions match
+    // (Checked in debug assertions).
+    #[inline]
+    fn broadcast_assume<E>(&self, dim: E) -> ArrayView<A, E>
+        where E: Dimension,
+    {
+        let dim = dim.into_dimension();
+        debug_assert_eq!(self.shape(), dim.slice());
+        let ptr = self.ptr;
+        let mut strides = dim.clone();
+        strides.slice_mut().copy_from_slice(self.strides.slice());
+        unsafe {
+            ArrayView::new_(ptr, dim, strides)
+        }
+    }
+
     fn raw_strides(&self) -> D {
         self.strides.clone()
     }