Use IntoIterator for select_rows and select_columns argument.

This commit is contained in:
sebcrozet 2018-12-09 16:56:27 +01:00
parent d1391592a0
commit da8dc6c4bd
1 changed files with 13 additions and 7 deletions

View File

@ -35,15 +35,18 @@ impl<N: Scalar + Zero, R: Dim, C: Dim, S: Storage<N, R, C>> Matrix<N, R, C, S> {
} }
/// Creates a new matrix by extracting the given set of rows from `self`. /// Creates a new matrix by extracting the given set of rows from `self`.
pub fn select_rows(&self, irows: impl ExactSizeIterator<Item = usize> + Clone) -> MatrixMN<N, Dynamic, C> pub fn select_rows<'a, I>(&self, irows: I) -> MatrixMN<N, Dynamic, C>
where DefaultAllocator: Allocator<N, Dynamic, C> { where I: IntoIterator<Item = &'a usize>,
I::IntoIter: ExactSizeIterator + Clone,
DefaultAllocator: Allocator<N, Dynamic, C> {
let mut irows = irows.into_iter();
let ncols = self.data.shape().1; let ncols = self.data.shape().1;
let mut res = unsafe { MatrixMN::new_uninitialized_generic(Dynamic::new(irows.len()), ncols) }; let mut res = unsafe { MatrixMN::new_uninitialized_generic(Dynamic::new(irows.len()), ncols) };
// First, check that all the indices from irows are valid. // First, check that all the indices from irows are valid.
// This will allow us to use unchecked access in the inner loop. // This will allow us to use unchecked access in the inner loop.
for i in irows.clone() { for i in irows.clone() {
assert!(i < self.nrows(), "Row index out of bounds.") assert!(*i < self.nrows(), "Row index out of bounds.")
} }
for j in 0..ncols.value() { for j in 0..ncols.value() {
@ -53,7 +56,7 @@ impl<N: Scalar + Zero, R: Dim, C: Dim, S: Storage<N, R, C>> Matrix<N, R, C, S> {
for (destination, source) in irows.clone().enumerate() { for (destination, source) in irows.clone().enumerate() {
unsafe { unsafe {
*res.vget_unchecked_mut(destination) = *src.vget_unchecked(source) *res.vget_unchecked_mut(destination) = *src.vget_unchecked(*source)
} }
} }
} }
@ -62,13 +65,16 @@ impl<N: Scalar + Zero, R: Dim, C: Dim, S: Storage<N, R, C>> Matrix<N, R, C, S> {
} }
/// Creates a new matrix by extracting the given set of columns from `self`. /// Creates a new matrix by extracting the given set of columns from `self`.
pub fn select_columns(&self, icols: impl ExactSizeIterator<Item = usize>) -> MatrixMN<N, R, Dynamic> pub fn select_columns<'a, I>(&self, icols: I) -> MatrixMN<N, R, Dynamic>
where DefaultAllocator: Allocator<N, R, Dynamic> { where I: IntoIterator<Item = &'a usize>,
I::IntoIter: ExactSizeIterator,
DefaultAllocator: Allocator<N, R, Dynamic> {
let mut icols = icols.into_iter();
let nrows = self.data.shape().0; let nrows = self.data.shape().0;
let mut res = unsafe { MatrixMN::new_uninitialized_generic(nrows, Dynamic::new(icols.len())) }; let mut res = unsafe { MatrixMN::new_uninitialized_generic(nrows, Dynamic::new(icols.len())) };
for (destination, source) in icols.enumerate() { for (destination, source) in icols.enumerate() {
res.column_mut(destination).copy_from(&self.column(source)) res.column_mut(destination).copy_from(&self.column(*source))
} }
res res