nalgebra/src/base/iter.rs

402 lines
12 KiB
Rust
Raw Normal View History

2017-02-13 01:17:09 +08:00
//! Matrix iterators.
2020-11-20 05:00:28 +08:00
use std::iter::FusedIterator;
use std::marker::PhantomData;
use std::mem;
2019-03-23 21:29:07 +08:00
use crate::base::dimension::{Dim, U1};
use crate::base::storage::{RawStorage, RawStorageMut};
use crate::base::{Matrix, MatrixSlice, MatrixSliceMut, Scalar};
macro_rules! iterator {
(struct $Name:ident for $Storage:ident.$ptr: ident -> $Ptr:ty, $Ref:ty, $SRef: ty) => {
/// An iterator through a dense matrix with arbitrary strides matrix.
#[derive(Debug)]
2021-07-14 17:25:16 +08:00
pub struct $Name<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> {
2018-10-22 13:00:10 +08:00
ptr: $Ptr,
inner_ptr: $Ptr,
inner_end: $Ptr,
2018-10-22 13:00:10 +08:00
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)>,
}
2020-11-15 23:57:49 +08:00
// 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).
2021-07-14 17:25:16 +08:00
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.
2021-04-11 17:00:38 +08:00
pub fn new(storage: $SRef) -> $Name<'a, T, R, C, S> {
2018-10-22 13:00:10 +08:00
let shape = storage.shape();
let strides = storage.strides();
let inner_offset = shape.0.value() * strides.0.value();
let size = shape.0.value() * shape.1.value();
let ptr = storage.$ptr();
// If we have a size of 0, 'ptr' must be
// dangling. Howver, 'inner_offset' might
// not be zero if only one dimension is zero, so
// we don't want to call 'offset'.
// This pointer will never actually get used
// if our size is '0', so it's fine to use
// 'ptr' for both the start and end.
let inner_end = if size == 0 {
ptr
} else {
// Safety:
// If 'size' is non-zero, we know that 'ptr'
// is not dangling, and 'inner_offset' must lie
// within the allocation
unsafe { ptr.add(inner_offset) }
};
$Name {
ptr,
inner_ptr: ptr,
inner_end,
2018-10-22 13:00:10 +08:00
size: shape.0.value() * shape.1.value(),
strides,
2018-10-22 13:00:10 +08:00
_phantoms: PhantomData,
}
}
}
2021-07-14 17:25:16 +08:00
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]
fn next(&mut self) -> Option<$Ref> {
unsafe {
if self.size == 0 {
None
} else {
self.size -= 1;
// Jump to the next outer dimension if needed.
if self.ptr == self.inner_end {
let stride = self.strides.1.value() as isize;
// This might go past the end of the allocation,
// depending on the value of 'size'. We use
// `wrapping_offset` to avoid UB
self.inner_end = self.ptr.wrapping_offset(stride);
// This will always be in bounds, since
// we're going to dereference it
self.ptr = self.inner_ptr.offset(stride);
self.inner_ptr = self.ptr;
}
// Go to the next element.
let old = self.ptr;
// Don't offset `self.ptr` for the last element,
// as this will be out of bounds. Iteration is done
// at this point (the next call to `next` will return `None`)
// so this is not observable.
if self.size != 0 {
let stride = self.strides.0.value();
self.ptr = self.ptr.add(stride);
}
2021-06-18 15:45:37 +08:00
// We want either `& *last` or `&mut *last` here, depending
// on the mutability of `$Ref`.
#[allow(clippy::transmute_ptr_to_ref)]
Some(mem::transmute(old))
}
}
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(self.size, Some(self.size))
}
#[inline]
fn count(self) -> usize {
self.size_hint().0
}
}
2021-07-14 17:25:16 +08:00
impl<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> DoubleEndedIterator
2021-04-11 17:00:38 +08:00
for $Name<'a, T, R, C, S>
{
#[inline]
fn next_back(&mut self) -> Option<$Ref> {
2020-11-20 03:27:59 +08:00
unsafe {
if self.size == 0 {
None
} else {
// Pre-decrement `size` such that it now counts to the
// element we want to return.
self.size -= 1;
2020-11-20 08:30:59 +08:00
// Fetch strides
let inner_stride = self.strides.0.value();
let outer_stride = self.strides.1.value();
2020-11-20 04:43:29 +08:00
// Compute number of rows
2020-11-20 08:30:59 +08:00
// Division should be exact
let inner_raw_size = self.inner_end.offset_from(self.inner_ptr) as usize;
let inner_size = inner_raw_size / inner_stride;
2020-11-20 04:43:29 +08:00
// Compute rows and cols remaining
let outer_remaining = self.size / inner_size;
let inner_remaining = self.size % inner_size;
2020-11-20 03:27:59 +08:00
// Compute pointer to last element
2021-06-18 15:45:37 +08:00
let last = self
.ptr
.add((outer_remaining * outer_stride + inner_remaining * inner_stride));
2020-11-20 03:27:59 +08:00
// We want either `& *last` or `&mut *last` here, depending
// on the mutability of `$Ref`.
2021-06-18 15:45:37 +08:00
#[allow(clippy::transmute_ptr_to_ref)]
2020-11-20 03:27:59 +08:00
Some(mem::transmute(last))
}
}
}
}
2020-11-20 05:00:28 +08:00
2021-07-14 17:25:16 +08:00
impl<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> ExactSizeIterator
2021-04-11 17:00:38 +08:00
for $Name<'a, T, R, C, S>
2020-11-20 05:00:28 +08:00
{
#[inline]
fn len(&self) -> usize {
self.size
}
}
2021-07-14 17:25:16 +08:00
impl<'a, T, R: Dim, C: Dim, S: 'a + $Storage<T, R, C>> FusedIterator
2021-04-11 17:00:38 +08:00
for $Name<'a, T, R, C, S>
2020-11-20 05:00:28 +08:00
{
}
2018-10-22 13:00:10 +08:00
};
}
iterator!(struct MatrixIter for RawStorage.ptr -> *const T, &'a T, &'a S);
iterator!(struct MatrixIterMut for RawStorageMut.ptr_mut -> *mut T, &'a mut T, &'a mut S);
2018-12-09 22:21:05 +08:00
/*
*
* Row iterators.
*
*/
2021-07-28 07:18:29 +08:00
#[derive(Clone, Debug)]
/// An iterator through the rows of a matrix.
pub struct RowIter<'a, T, R: Dim, C: Dim, S: RawStorage<T, R, C>> {
2021-04-11 17:00:38 +08:00
mat: &'a Matrix<T, R, C, S>,
2020-04-06 00:49:48 +08:00
curr: usize,
2018-12-09 22:21:05 +08:00
}
impl<'a, T, R: Dim, C: Dim, S: 'a + RawStorage<T, R, C>> RowIter<'a, T, R, C, S> {
2021-04-11 17:00:38 +08:00
pub(crate) fn new(mat: &'a Matrix<T, R, C, S>) -> Self {
2020-04-06 00:49:48 +08:00
RowIter { mat, curr: 0 }
2018-12-09 22:21:05 +08:00
}
}
impl<'a, T, R: Dim, C: Dim, S: 'a + RawStorage<T, R, C>> Iterator for RowIter<'a, T, R, C, S> {
2021-04-11 17:00:38 +08:00
type Item = MatrixSlice<'a, T, U1, C, S::RStride, S::CStride>;
2018-12-09 22:21:05 +08:00
#[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>) {
2020-04-06 00:49:48 +08:00
(
self.mat.nrows() - self.curr,
Some(self.mat.nrows() - self.curr),
)
2018-12-09 22:21:05 +08:00
}
#[inline]
fn count(self) -> usize {
self.mat.nrows() - self.curr
}
}
impl<'a, T: Scalar, R: Dim, C: Dim, S: 'a + RawStorage<T, R, C>> ExactSizeIterator
2021-04-11 17:00:38 +08:00
for RowIter<'a, T, R, C, S>
2020-04-06 00:49:48 +08:00
{
2018-12-09 22:21:05 +08:00
#[inline]
fn len(&self) -> usize {
self.mat.nrows() - self.curr
}
}
/// An iterator through the mutable rows of a matrix.
#[derive(Debug)]
pub struct RowIterMut<'a, T, R: Dim, C: Dim, S: RawStorageMut<T, R, C>> {
2021-04-11 17:00:38 +08:00
mat: *mut Matrix<T, R, C, S>,
2018-12-09 22:21:05 +08:00
curr: usize,
2021-04-11 17:00:38 +08:00
phantom: PhantomData<&'a mut Matrix<T, R, C, S>>,
2018-12-09 22:21:05 +08:00
}
impl<'a, T, R: Dim, C: Dim, S: 'a + RawStorageMut<T, R, C>> RowIterMut<'a, T, R, C, S> {
2021-04-11 17:00:38 +08:00
pub(crate) fn new(mat: &'a mut Matrix<T, R, C, S>) -> Self {
2018-12-09 22:21:05 +08:00
RowIterMut {
mat,
curr: 0,
2020-04-06 00:49:48 +08:00
phantom: PhantomData,
2018-12-09 22:21:05 +08:00
}
}
fn nrows(&self) -> usize {
2020-04-06 00:49:48 +08:00
unsafe { (*self.mat).nrows() }
2018-12-09 22:21:05 +08:00
}
}
impl<'a, T, R: Dim, C: Dim, S: 'a + RawStorageMut<T, R, C>> Iterator
for RowIterMut<'a, T, R, C, S>
{
2021-04-11 17:00:38 +08:00
type Item = MatrixSliceMut<'a, T, U1, C, S::RStride, S::CStride>;
2018-12-09 22:21:05 +08:00
#[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, T: Scalar, R: Dim, C: Dim, S: 'a + RawStorageMut<T, R, C>> ExactSizeIterator
2021-04-11 17:00:38 +08:00
for RowIterMut<'a, T, R, C, S>
2020-04-06 00:49:48 +08:00
{
2018-12-09 22:21:05 +08:00
#[inline]
fn len(&self) -> usize {
self.nrows() - self.curr
}
}
/*
*
* Column iterators.
*
*/
2021-07-28 07:18:29 +08:00
#[derive(Clone, Debug)]
/// An iterator through the columns of a matrix.
pub struct ColumnIter<'a, T, R: Dim, C: Dim, S: RawStorage<T, R, C>> {
2021-04-11 17:00:38 +08:00
mat: &'a Matrix<T, R, C, S>,
2020-04-06 00:49:48 +08:00
curr: usize,
2018-12-09 22:21:05 +08:00
}
impl<'a, T, R: Dim, C: Dim, S: 'a + RawStorage<T, R, C>> ColumnIter<'a, T, R, C, S> {
2021-04-11 17:00:38 +08:00
pub(crate) fn new(mat: &'a Matrix<T, R, C, S>) -> Self {
2020-04-06 00:49:48 +08:00
ColumnIter { mat, curr: 0 }
2018-12-09 22:21:05 +08:00
}
}
impl<'a, T, R: Dim, C: Dim, S: 'a + RawStorage<T, R, C>> Iterator for ColumnIter<'a, T, R, C, S> {
2021-04-11 17:00:38 +08:00
type Item = MatrixSlice<'a, T, R, U1, S::RStride, S::CStride>;
2018-12-09 22:21:05 +08:00
#[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>) {
2020-04-06 00:49:48 +08:00
(
self.mat.ncols() - self.curr,
Some(self.mat.ncols() - self.curr),
)
2018-12-09 22:21:05 +08:00
}
#[inline]
fn count(self) -> usize {
self.mat.ncols() - self.curr
}
}
impl<'a, T: Scalar, R: Dim, C: Dim, S: 'a + RawStorage<T, R, C>> ExactSizeIterator
2021-04-11 17:00:38 +08:00
for ColumnIter<'a, T, R, C, S>
2020-04-06 00:49:48 +08:00
{
2018-12-09 22:21:05 +08:00
#[inline]
fn len(&self) -> usize {
self.mat.ncols() - self.curr
}
}
/// An iterator through the mutable columns of a matrix.
#[derive(Debug)]
pub struct ColumnIterMut<'a, T, R: Dim, C: Dim, S: RawStorageMut<T, R, C>> {
2021-04-11 17:00:38 +08:00
mat: *mut Matrix<T, R, C, S>,
2018-12-09 22:21:05 +08:00
curr: usize,
2021-04-11 17:00:38 +08:00
phantom: PhantomData<&'a mut Matrix<T, R, C, S>>,
2018-12-09 22:21:05 +08:00
}
impl<'a, T, R: Dim, C: Dim, S: 'a + RawStorageMut<T, R, C>> ColumnIterMut<'a, T, R, C, S> {
2021-04-11 17:00:38 +08:00
pub(crate) fn new(mat: &'a mut Matrix<T, R, C, S>) -> Self {
2018-12-09 22:21:05 +08:00
ColumnIterMut {
mat,
curr: 0,
2020-04-06 00:49:48 +08:00
phantom: PhantomData,
2018-12-09 22:21:05 +08:00
}
}
fn ncols(&self) -> usize {
2020-04-06 00:49:48 +08:00
unsafe { (*self.mat).ncols() }
2018-12-09 22:21:05 +08:00
}
}
impl<'a, T, R: Dim, C: Dim, S: 'a + RawStorageMut<T, R, C>> Iterator
2021-04-11 17:00:38 +08:00
for ColumnIterMut<'a, T, R, C, S>
2020-04-06 00:49:48 +08:00
{
2021-04-11 17:00:38 +08:00
type Item = MatrixSliceMut<'a, T, R, U1, S::RStride, S::CStride>;
2018-12-09 22:21:05 +08:00
#[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, T: Scalar, R: Dim, C: Dim, S: 'a + RawStorageMut<T, R, C>> ExactSizeIterator
2021-04-11 17:00:38 +08:00
for ColumnIterMut<'a, T, R, C, S>
2020-04-06 00:49:48 +08:00
{
2018-12-09 22:21:05 +08:00
#[inline]
fn len(&self) -> usize {
self.ncols() - self.curr
}
}