srenblad/nalgebra
1
0
Fork 0
forked from M-Labs/nalgebra
Code Pull Requests Activity

Add row and column iterators.

Browse Source
This commit is contained in:
sebcrozet 2018-12-09 15:21:05 +01:00
parent bba1f48e81
commit 0d2c1be8da
2 changed files with 246 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

225
src/base/iter.rs
View File

@ -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<N, R, C>> {
mat: &'a Matrix<N, R, C, S>,
curr: usize
}
impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + Storage<N, R, C>> RowIter<'a, N, R, C, S> {
pub fn new(mat: &'a Matrix<N, R, C, S>) -> Self {
RowIter {
mat, curr: 0
}
}
}
impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + Storage<N, R, C>> 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<Self::Item> {
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<usize>) {
(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<N, R, C>> 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<N, R, C>> {
mat: *mut Matrix<N, R, C, S>,
curr: usize,
phantom: PhantomData<&'a Matrix<N, R, C, S>>
}
impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + StorageMut<N, R, C>> RowIterMut<'a, N, R, C, S> {
pub fn new(mat: &'a mut Matrix<N, R, C, S>) -> 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<N, R, C>> 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<Self::Item> {
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<usize>) {
(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<N, R, C>> 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<N, R, C>> {
mat: &'a Matrix<N, R, C, S>,
curr: usize
}
impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + Storage<N, R, C>> ColumnIter<'a, N, R, C, S> {
pub fn new(mat: &'a Matrix<N, R, C, S>) -> Self {
ColumnIter {
mat, curr: 0
}
}
}
impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + Storage<N, R, C>> 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<Self::Item> {
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<usize>) {
(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<N, R, C>> 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<N, R, C>> {
mat: *mut Matrix<N, R, C, S>,
curr: usize,
phantom: PhantomData<&'a Matrix<N, R, C, S>>
}
impl<'a, N: Scalar, R: Dim, C: Dim, S: 'a + StorageMut<N, R, C>> ColumnIterMut<'a, N, R, C, S> {
pub fn new(mat: &'a mut Matrix<N, R, C, S>) -> 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<N, R, C>> 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<Self::Item> {
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<usize>) {
(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<N, R, C>> ExactSizeIterator for ColumnIterMut<'a, N, R, C, S> {
#[inline]
fn len(&self) -> usize {
self.ncols() - self.curr
}
}

24
src/base/matrix.rs
View File

@ -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<N: Scalar, R: Dim, C: Dim, S: Storage<N, R, C>> Matrix<N, R, C, S> {
MatrixIter::new(&self.data)
}
#[inline]
pub fn row_iter(&self) -> RowIter<N, R, C, S> {
RowIter::new(self)
}
#[inline]
pub fn column_iter(&self) -> ColumnIter<N, R, C, S> {
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<N: Scalar, R: Dim, C: Dim, S: StorageMut<N, R, C>> Matrix<N, R, C, S> {
MatrixIterMut::new(&mut self.data)
}
/// Mutably iterates through this matrix rows.
#[inline]
pub fn row_iter_mut(&mut self) -> RowIterMut<N, R, C, S> {
RowIterMut::new(self)
}
/// Mutably iterates through this matrix columns.
#[inline]
pub fn column_iter_mut(&mut self) -> ColumnIterMut<N, R, C, S> {
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)) {