From 0d2c1be8da888bcec291478427ed278be347fcd6 Mon Sep 17 00:00:00 2001 From: sebcrozet Date: Sun, 9 Dec 2018 15:21:05 +0100 Subject: [PATCH] Add row and column iterators. --- src/base/iter.rs | 225 ++++++++++++++++++++++++++++++++++++++++++++- src/base/matrix.rs | 24 ++++- 2 files changed, 246 insertions(+), 3 deletions(-) diff --git a/src/base/iter.rs b/src/base/iter.rs index fe8fd4f6..2b5a12d4 100644 --- a/src/base/iter.rs +++ b/src/base/iter.rs @@ -3,9 +3,9 @@ use std::marker::PhantomData; use std::mem; -use base::dimension::Dim; +use base::dimension::{Dim, U1}; use base::storage::{Storage, StorageMut}; -use base::Scalar; +use base::{Scalar, Matrix, MatrixSlice, MatrixSliceMut, Dynamic}; macro_rules! iterator { (struct $Name:ident for $Storage:ident.$ptr: ident -> $Ptr:ty, $Ref:ty, $SRef: ty) => { @@ -96,3 +96,224 @@ macro_rules! iterator { iterator!(struct MatrixIter for Storage.ptr -> *const N, &'a N, &'a S); iterator!(struct MatrixIterMut for StorageMut.ptr_mut -> *mut N, &'a mut N, &'a mut S); + + +/* + * + * Row iterators. + * + */ +#[derive(Clone)] +pub struct RowIter<'a, N: Scalar, R: Dim, C: Dim, S: Storage> { + mat: &'a Matrix, + curr: usize +} + +impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + Storage> RowIter<'a, N, R, C, S> { + pub fn new(mat: &'a Matrix) -> Self { + RowIter { + mat, curr: 0 + } + } +} + + +impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + Storage> Iterator for RowIter<'a, N, R, C, S> { + type Item = MatrixSlice<'a, N, U1, C, S::RStride, S::CStride>; + + #[inline] + fn next(&mut self) -> Option { + if self.curr < self.mat.nrows() { + let res = self.mat.row(self.curr); + self.curr += 1; + Some(res) + } else { + None + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + (self.mat.nrows() - self.curr, Some(self.mat.nrows() - self.curr)) + } + + #[inline] + fn count(self) -> usize { + self.mat.nrows() - self.curr + } +} + +impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + Storage> ExactSizeIterator for RowIter<'a, N, R, C, S> { + #[inline] + fn len(&self) -> usize { + self.mat.nrows() - self.curr + } +} + + + +pub struct RowIterMut<'a, N: Scalar, R: Dim, C: Dim, S: StorageMut> { + mat: *mut Matrix, + curr: usize, + phantom: PhantomData<&'a Matrix> +} + +impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + StorageMut> RowIterMut<'a, N, R, C, S> { + pub fn new(mat: &'a mut Matrix) -> Self { + RowIterMut { + mat, + curr: 0, + phantom: PhantomData + } + } + + fn nrows(&self) -> usize { + unsafe { + (*self.mat).nrows() + } + } +} + + +impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + StorageMut> Iterator for RowIterMut<'a, N, R, C, S> { + type Item = MatrixSliceMut<'a, N, U1, C, S::RStride, S::CStride>; + + #[inline] + fn next(&mut self) -> Option { + if self.curr < self.nrows() { + let res = unsafe { (*self.mat).row_mut(self.curr) }; + self.curr += 1; + Some(res) + } else { + None + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + (self.nrows() - self.curr, Some(self.nrows() - self.curr)) + } + + #[inline] + fn count(self) -> usize { + self.nrows() - self.curr + } +} + +impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + StorageMut> ExactSizeIterator for RowIterMut<'a, N, R, C, S> { + #[inline] + fn len(&self) -> usize { + self.nrows() - self.curr + } +} + + +/* + * + * Column iterators. + * + */ +#[derive(Clone)] +pub struct ColumnIter<'a, N: Scalar, R: Dim, C: Dim, S: Storage> { + mat: &'a Matrix, + curr: usize +} + +impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + Storage> ColumnIter<'a, N, R, C, S> { + pub fn new(mat: &'a Matrix) -> Self { + ColumnIter { + mat, curr: 0 + } + } +} + + +impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + Storage> Iterator for ColumnIter<'a, N, R, C, S> { + type Item = MatrixSlice<'a, N, R, U1, S::RStride, S::CStride>; + + #[inline] + fn next(&mut self) -> Option { + if self.curr < self.mat.ncols() { + let res = self.mat.column(self.curr); + self.curr += 1; + Some(res) + } else { + None + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + (self.mat.ncols() - self.curr, Some(self.mat.ncols() - self.curr)) + } + + #[inline] + fn count(self) -> usize { + self.mat.ncols() - self.curr + } +} + +impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + Storage> ExactSizeIterator for ColumnIter<'a, N, R, C, S> { + #[inline] + fn len(&self) -> usize { + self.mat.ncols() - self.curr + } +} + + + +pub struct ColumnIterMut<'a, N: Scalar, R: Dim, C: Dim, S: StorageMut> { + mat: *mut Matrix, + curr: usize, + phantom: PhantomData<&'a Matrix> +} + +impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + StorageMut> ColumnIterMut<'a, N, R, C, S> { + pub fn new(mat: &'a mut Matrix) -> Self { + ColumnIterMut { + mat, + curr: 0, + phantom: PhantomData + } + } + + fn ncols(&self) -> usize { + unsafe { + (*self.mat).ncols() + } + } +} + + +impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + StorageMut> Iterator for ColumnIterMut<'a, N, R, C, S> { + type Item = MatrixSliceMut<'a, N, R, U1, S::RStride, S::CStride>; + + #[inline] + fn next(&mut self) -> Option { + if self.curr < self.ncols() { + let res = unsafe { (*self.mat).column_mut(self.curr) }; + self.curr += 1; + Some(res) + } else { + None + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + (self.ncols() - self.curr, Some(self.ncols() - self.curr)) + } + + #[inline] + fn count(self) -> usize { + self.ncols() - self.curr + } +} + +impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + StorageMut> ExactSizeIterator for ColumnIterMut<'a, N, R, C, S> { + #[inline] + fn len(&self) -> usize { + self.ncols() - self.curr + } +} + diff --git a/src/base/matrix.rs b/src/base/matrix.rs index d9c03b54..625438d5 100644 --- a/src/base/matrix.rs +++ b/src/base/matrix.rs @@ -22,7 +22,7 @@ use alga::general::{ClosedAdd, ClosedMul, ClosedSub, Real, Ring}; use base::allocator::{Allocator, SameShapeAllocator, SameShapeC, SameShapeR}; use base::constraint::{DimEq, SameNumberOfColumns, SameNumberOfRows, ShapeConstraint}; use base::dimension::{Dim, DimAdd, DimSum, IsNotStaticOne, U1, U2, U3}; -use base::iter::{MatrixIter, MatrixIterMut}; +use base::iter::{MatrixIter, MatrixIterMut, RowIter, RowIterMut, ColumnIter, ColumnIterMut}; use base::storage::{ ContiguousStorage, ContiguousStorageMut, Owned, SameShapeStorage, Storage, StorageMut, }; @@ -247,6 +247,16 @@ impl> Matrix { MatrixIter::new(&self.data) } + #[inline] + pub fn row_iter(&self) -> RowIter { + RowIter::new(self) + } + + #[inline] + pub fn column_iter(&self) -> ColumnIter { + ColumnIter::new(self) + } + /// Computes the row and column coordinates of the i-th element of this matrix seen as a /// vector. #[inline] @@ -591,6 +601,18 @@ impl> Matrix { MatrixIterMut::new(&mut self.data) } + /// Mutably iterates through this matrix rows. + #[inline] + pub fn row_iter_mut(&mut self) -> RowIterMut { + RowIterMut::new(self) + } + + /// Mutably iterates through this matrix columns. + #[inline] + pub fn column_iter_mut(&mut self) -> ColumnIterMut { + ColumnIterMut::new(self) + } + /// Swaps two entries without bound-checking. #[inline] pub unsafe fn swap_unchecked(&mut self, row_cols1: (usize, usize), row_cols2: (usize, usize)) {