Merge pull request #1315 from yotamofek/owned-view-iter

Allow creating matrix iter with an owned view

src/base/conversion.rs

@@ -98,6 +98,18 @@ impl<'a, T: Scalar, R: Dim, C: Dim, S: RawStorage<T, R, C>> IntoIterator
     }
 }
 
+impl<'a, T: Scalar, R: Dim, C: Dim, RStride: Dim, CStride: Dim> IntoIterator
+    for Matrix<T, R, C, ViewStorage<'a, T, R, C, RStride, CStride>>
+{
+    type Item = &'a T;
+    type IntoIter = MatrixIter<'a, T, R, C, ViewStorage<'a, T, R, C, RStride, CStride>>;
+
+    #[inline]
+    fn into_iter(self) -> Self::IntoIter {
+        MatrixIter::new_owned(self.data)
+    }
+}
+
 impl<'a, T: Scalar, R: Dim, C: Dim, S: RawStorageMut<T, R, C>> IntoIterator
     for &'a mut Matrix<T, R, C, S>
 {
@@ -110,6 +122,18 @@ impl<'a, T: Scalar, R: Dim, C: Dim, S: RawStorageMut<T, R, C>> IntoIterator
     }
 }
 
+impl<'a, T: Scalar, R: Dim, C: Dim, RStride: Dim, CStride: Dim> IntoIterator
+    for Matrix<T, R, C, ViewStorageMut<'a, T, R, C, RStride, CStride>>
+{
+    type Item = &'a mut T;
+    type IntoIter = MatrixIterMut<'a, T, R, C, ViewStorageMut<'a, T, R, C, RStride, CStride>>;
+
+    #[inline]
+    fn into_iter(self) -> Self::IntoIter {
+        MatrixIterMut::new_owned_mut(self.data)
+    }
+}
+
 impl<T: Scalar, const D: usize> From<[T; D]> for SVector<T, D> {
     #[inline]
     fn from(arr: [T; D]) -> Self {

src/base/iter.rs

@@ -12,26 +12,29 @@ use std::mem;
 
 use crate::base::dimension::{Dim, U1};
 use crate::base::storage::{RawStorage, RawStorageMut};
-use crate::base::{Matrix, MatrixView, MatrixViewMut, Scalar};
+use crate::base::{Matrix, MatrixView, MatrixViewMut, Scalar, ViewStorage, ViewStorageMut};
+
+#[derive(Clone, Debug)]
+struct RawIter<Ptr, T, R: Dim, C: Dim, RStride: Dim, CStride: Dim> {
+    ptr: Ptr,
+    inner_ptr: Ptr,
+    inner_end: Ptr,
+    size: usize,
+    strides: (RStride, CStride),
+    _phantoms: PhantomData<(fn() -> T, R, C)>,
+}
 
 macro_rules! iterator {
     (struct $Name:ident for $Storage:ident.$ptr: ident -> $Ptr:ty, $Ref:ty, $SRef: ty, $($derives:ident),* $(,)?) => {
-        /// An iterator through a dense matrix with arbitrary strides matrix.
-        #[derive($($derives),*)]
-        pub struct $Name<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> {
-            ptr: $Ptr,
-            inner_ptr: $Ptr,
-            inner_end: $Ptr,
-            size: usize, // We can't use an end pointer here because a stride might be zero.
-            strides: (S::RStride, S::CStride),
-            _phantoms: PhantomData<($Ref, R, C, S)>,
-        }
-
         // TODO: we need to specialize for the case where the matrix storage is owned (in which
         // case the iterator is trivial because it does not have any stride).
-        impl<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> $Name<'a, T, R, C, S> {
+        impl<T, R: Dim, C: Dim, RStride: Dim, CStride: Dim>
+            RawIter<$Ptr, T, R, C, RStride, CStride>
+        {
             /// Creates a new iterator for the given matrix storage.
-            pub fn new(storage: $SRef) -> $Name<'a, T, R, C, S> {
+            fn new<'a, S: $Storage<T, R, C, RStride = RStride, CStride = CStride>>(
+                storage: $SRef,
+            ) -> Self {
                 let shape = storage.shape();
                 let strides = storage.strides();
                 let inner_offset = shape.0.value() * strides.0.value();
@@ -55,7 +58,7 @@ macro_rules! iterator {
                     unsafe { ptr.add(inner_offset) }
                 };
 
-                $Name {
+                RawIter {
                     ptr,
                     inner_ptr: ptr,
                     inner_end,
@@ -66,11 +69,13 @@ macro_rules! iterator {
             }
         }
 
-        impl<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> Iterator for $Name<'a, T, R, C, S> {
+        impl<T, R: Dim, C: Dim, RStride: Dim, CStride: Dim> Iterator
-            type Item = $Ref;
+            for RawIter<$Ptr, T, R, C, RStride, CStride>
+        {
+            type Item = $Ptr;
 
             #[inline]
-            fn next(&mut self) -> Option<$Ref> {
+            fn next(&mut self) -> Option<Self::Item> {
                 unsafe {
                     if self.size == 0 {
                         None
@@ -102,10 +107,7 @@ macro_rules! iterator {
                             self.ptr = self.ptr.add(stride);
                         }
 
-                        // We want either `& *last` or `&mut *last` here, depending
+                        Some(old)
-                        // on the mutability of `$Ref`.
-                        #[allow(clippy::transmute_ptr_to_ref)]
-                        Some(mem::transmute(old))
                     }
                 }
             }
@@ -121,11 +123,11 @@ macro_rules! iterator {
             }
         }
 
-        impl<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> DoubleEndedIterator
+        impl<T, R: Dim, C: Dim, RStride: Dim, CStride: Dim> DoubleEndedIterator
-            for $Name<'a, T, R, C, S>
+            for RawIter<$Ptr, T, R, C, RStride, CStride>
         {
             #[inline]
-            fn next_back(&mut self) -> Option<$Ref> {
+            fn next_back(&mut self) -> Option<Self::Item> {
                 unsafe {
                     if self.size == 0 {
                         None
@@ -152,21 +154,85 @@ macro_rules! iterator {
                             .ptr
                             .add((outer_remaining * outer_stride + inner_remaining * inner_stride));
 
+                        Some(last)
+                    }
+                }
+            }
+        }
+
+        impl<T, R: Dim, C: Dim, RStride: Dim, CStride: Dim> ExactSizeIterator
+            for RawIter<$Ptr, T, R, C, RStride, CStride>
+        {
+            #[inline]
+            fn len(&self) -> usize {
+                self.size
+            }
+        }
+
+        impl<T, R: Dim, C: Dim, RStride: Dim, CStride: Dim> FusedIterator
+            for RawIter<$Ptr, T, R, C, RStride, CStride>
+        {
+        }
+
+        /// An iterator through a dense matrix with arbitrary strides matrix.
+        #[derive($($derives),*)]
+        pub struct $Name<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> {
+            inner: RawIter<$Ptr, T, R, C, S::RStride, S::CStride>,
+            _marker: PhantomData<$Ref>,
+        }
+
+        impl<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> $Name<'a, T, R, C, S> {
+            /// Creates a new iterator for the given matrix storage.
+            pub fn new(storage: $SRef) -> Self {
+                Self {
+                    inner: RawIter::<$Ptr, T, R, C, S::RStride, S::CStride>::new(storage),
+                    _marker: PhantomData,
+                }
+            }
+        }
+
+        impl<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> Iterator for $Name<'a, T, R, C, S> {
+            type Item = $Ref;
+
+            #[inline(always)]
+            fn next(&mut self) -> Option<Self::Item> {
                 // We want either `& *last` or `&mut *last` here, depending
                 // on the mutability of `$Ref`.
                 #[allow(clippy::transmute_ptr_to_ref)]
-                        Some(mem::transmute(last))
+                self.inner.next().map(|ptr| unsafe { mem::transmute(ptr) })
+            }
+
+            #[inline(always)]
+            fn size_hint(&self) -> (usize, Option<usize>) {
+                self.inner.size_hint()
+            }
+
+            #[inline(always)]
+            fn count(self) -> usize {
+                self.inner.count()
             }
         }
+
+        impl<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> DoubleEndedIterator
+            for $Name<'a, T, R, C, S>
+        {
+            #[inline(always)]
+            fn next_back(&mut self) -> Option<Self::Item> {
+                // We want either `& *last` or `&mut *last` here, depending
+                // on the mutability of `$Ref`.
+                #[allow(clippy::transmute_ptr_to_ref)]
+                self.inner
+                    .next_back()
+                    .map(|ptr| unsafe { mem::transmute(ptr) })
             }
         }
 
         impl<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> ExactSizeIterator
             for $Name<'a, T, R, C, S>
         {
-            #[inline]
+            #[inline(always)]
             fn len(&self) -> usize {
-                self.size
+                self.inner.len()
             }
         }
 
@@ -180,6 +246,30 @@ macro_rules! iterator {
 iterator!(struct MatrixIter for RawStorage.ptr -> *const T, &'a T, &'a S, Clone, Debug);
 iterator!(struct MatrixIterMut for RawStorageMut.ptr_mut -> *mut T, &'a mut T, &'a mut S, Debug);
 
+impl<'a, T, R: Dim, C: Dim, RStride: Dim, CStride: Dim>
+    MatrixIter<'a, T, R, C, ViewStorage<'a, T, R, C, RStride, CStride>>
+{
+    /// Creates a new iterator for the given matrix storage view.
+    pub fn new_owned(storage: ViewStorage<'a, T, R, C, RStride, CStride>) -> Self {
+        Self {
+            inner: RawIter::<*const T, T, R, C, RStride, CStride>::new(&storage),
+            _marker: PhantomData,
+        }
+    }
+}
+
+impl<'a, T, R: Dim, C: Dim, RStride: Dim, CStride: Dim>
+    MatrixIterMut<'a, T, R, C, ViewStorageMut<'a, T, R, C, RStride, CStride>>
+{
+    /// Creates a new iterator for the given matrix storage view.
+    pub fn new_owned_mut(mut storage: ViewStorageMut<'a, T, R, C, RStride, CStride>) -> Self {
+        Self {
+            inner: RawIter::<*mut T, T, R, C, RStride, CStride>::new(&mut storage),
+            _marker: PhantomData,
+        }
+    }
+}
+
 /*
  *
  * Row iterators.

tests/core/matrix.rs

@@ -1,18 +1,22 @@
+use na::iter::MatrixIter;
 use num::{One, Zero};
 use std::cmp::Ordering;
 
 use na::dimension::{U15, U8};
 use na::{
     self, Const, DMatrix, DVector, Matrix2, Matrix2x3, Matrix2x4, Matrix3, Matrix3x2, Matrix3x4,
-    Matrix4, Matrix4x3, Matrix4x5, Matrix5, Matrix6, OMatrix, RowVector3, RowVector4, RowVector5,
+    Matrix4, Matrix4x3, Matrix4x5, Matrix5, Matrix6, MatrixView2x3, MatrixViewMut2x3, OMatrix,
-    Vector1, Vector2, Vector3, Vector4, Vector5, Vector6,
+    RowVector3, RowVector4, RowVector5, Vector1, Vector2, Vector3, Vector4, Vector5, Vector6,
 };
 
 #[test]
 fn iter() {
     let a = Matrix2x3::new(1.0, 2.0, 3.0, 4.0, 5.0, 6.0);
+    let view: MatrixView2x3<_> = (&a).into();
 
-    let mut it = a.iter();
+    fn test<'a, F: Fn() -> I, I: Iterator<Item = &'a f64> + DoubleEndedIterator>(it: F) {
+        {
+            let mut it = it();
             assert_eq!(*it.next().unwrap(), 1.0);
             assert_eq!(*it.next().unwrap(), 4.0);
             assert_eq!(*it.next().unwrap(), 2.0);
@@ -20,8 +24,10 @@ fn iter() {
             assert_eq!(*it.next().unwrap(), 3.0);
             assert_eq!(*it.next().unwrap(), 6.0);
             assert!(it.next().is_none());
+        }
 
-    let mut it = a.iter();
+        {
+            let mut it = it();
             assert_eq!(*it.next().unwrap(), 1.0);
             assert_eq!(*it.next_back().unwrap(), 6.0);
             assert_eq!(*it.next_back().unwrap(), 3.0);
@@ -29,8 +35,9 @@ fn iter() {
             assert_eq!(*it.next().unwrap(), 4.0);
             assert_eq!(*it.next().unwrap(), 2.0);
             assert!(it.next().is_none());
-
+        }
-    let mut it = a.iter().rev();
+        {
+            let mut it = it().rev();
             assert_eq!(*it.next().unwrap(), 6.0);
             assert_eq!(*it.next().unwrap(), 3.0);
             assert_eq!(*it.next().unwrap(), 5.0);
@@ -38,43 +45,82 @@ fn iter() {
             assert_eq!(*it.next().unwrap(), 4.0);
             assert_eq!(*it.next().unwrap(), 1.0);
             assert!(it.next().is_none());
+        }
+    }
+
+    test(|| a.iter());
+    test(|| view.into_iter());
 
     let row = a.row(0);
-    let mut it = row.iter();
+    let row_test = |mut it: MatrixIter<_, _, _, _>| {
         assert_eq!(*it.next().unwrap(), 1.0);
         assert_eq!(*it.next().unwrap(), 2.0);
         assert_eq!(*it.next().unwrap(), 3.0);
         assert!(it.next().is_none());
+    };
+    row_test(row.iter());
+    row_test(row.into_iter());
 
     let row = a.row(1);
-    let mut it = row.iter();
+    let row_test = |mut it: MatrixIter<_, _, _, _>| {
         assert_eq!(*it.next().unwrap(), 4.0);
         assert_eq!(*it.next().unwrap(), 5.0);
         assert_eq!(*it.next().unwrap(), 6.0);
         assert!(it.next().is_none());
+    };
+    row_test(row.iter());
+    row_test(row.into_iter());
 
     let m22 = row.column(1);
-    let mut it = m22.iter();
+    let m22_test = |mut it: MatrixIter<_, _, _, _>| {
         assert_eq!(*it.next().unwrap(), 5.0);
         assert!(it.next().is_none());
+    };
+    m22_test(m22.iter());
+    m22_test(m22.into_iter());
 
     let col = a.column(0);
-    let mut it = col.iter();
+    let col_test = |mut it: MatrixIter<_, _, _, _>| {
         assert_eq!(*it.next().unwrap(), 1.0);
         assert_eq!(*it.next().unwrap(), 4.0);
         assert!(it.next().is_none());
+    };
+    col_test(col.iter());
+    col_test(col.into_iter());
 
     let col = a.column(1);
-    let mut it = col.iter();
+    let col_test = |mut it: MatrixIter<_, _, _, _>| {
         assert_eq!(*it.next().unwrap(), 2.0);
         assert_eq!(*it.next().unwrap(), 5.0);
         assert!(it.next().is_none());
+    };
+    col_test(col.iter());
+    col_test(col.into_iter());
 
     let col = a.column(2);
-    let mut it = col.iter();
+    let col_test = |mut it: MatrixIter<_, _, _, _>| {
         assert_eq!(*it.next().unwrap(), 3.0);
         assert_eq!(*it.next().unwrap(), 6.0);
         assert!(it.next().is_none());
+    };
+    col_test(col.iter());
+    col_test(col.into_iter());
+}
+
+#[test]
+fn iter_mut() {
+    let mut a = Matrix2x3::new(1.0, 2.0, 3.0, 4.0, 5.0, 6.0);
+
+    for v in a.iter_mut() {
+        *v *= 2.0;
+    }
+    assert_eq!(a, Matrix2x3::new(2.0, 4.0, 6.0, 8.0, 10.0, 12.0));
+
+    let view: MatrixViewMut2x3<_> = MatrixViewMut2x3::from(&mut a);
+    for v in view.into_iter() {
+        *v *= 2.0;
+    }
+    assert_eq!(a, Matrix2x3::new(4.0, 8.0, 12.0, 16.0, 20.0, 24.0));
 }
 
 #[test]