nalgebra/src/core/matrix_slice.rs

use std::marker::PhantomData;

use core::{Scalar, Matrix};
use core::dimension::{Dim, DimName, Dynamic, DimMul, DimProd, U1};
use core::iter::MatrixIter;
use core::storage::{Storage, StorageMut, Owned};
use core::allocator::Allocator;

macro_rules! slice_storage_impl(
    ($doc: expr; $Storage: ident as $SRef: ty; $T: ident.$get_addr: ident ($Ptr: ty as $Ref: ty)) => {
        #[doc = $doc]
        pub struct $T<'a, N: Scalar, R: Dim, C: Dim, RStride: Dim, CStride: Dim, Alloc> {
            ptr:       $Ptr,
            shape:     (R, C),
            strides:   (RStride, CStride),
            _phantoms: PhantomData<($Ref, Alloc)>,
        }

        // Dynamic and () are arbitrary. It's just to be able to call the constructors with
        // `Slice::`
        impl<'a, N: Scalar, R: Dim, C: Dim> $T<'a, N, R, C, Dynamic, Dynamic, ()> {
            /// Create a new matrix slice without bound checking.
            #[inline]
            pub unsafe fn new_unchecked<RStor, CStor, S>(storage: $SRef, start: (usize, usize), shape: (R, C))
                -> $T<'a, N, R, C, S::RStride, S::CStride, S::Alloc>
                where RStor: Dim,
                      CStor: Dim,
                      S:     $Storage<N, RStor, CStor> {

                let strides = storage.strides();
                $T::new_with_strides_unchecked(storage, start, shape, strides)
            }

            /// Create a new matrix slice without bound checking.
            #[inline]
            pub unsafe fn new_with_strides_unchecked<S, RStor, CStor, RStride, CStride>(storage: $SRef,
                                                                                        start:   (usize, usize),
                                                                                        shape:   (R, C),
                                                                                        strides: (RStride, CStride))
                -> $T<'a, N, R, C, RStride, CStride, S::Alloc>
                where RStor: Dim,
                      CStor: Dim,
                      S: $Storage<N, RStor, CStor>,
                      RStride: Dim,
                      CStride: Dim {

                $T {
                    ptr:       storage.$get_addr(start.0, start.1),
                    shape:     shape,
                    strides:   (strides.0, strides.1),
                    _phantoms: PhantomData
                }
            }
        }
    }
);

slice_storage_impl!("A matrix data storage for a matrix slice. Only contains an internal reference \
                     to another matrix data storage.";
    Storage as &'a S; SliceStorage.get_address_unchecked(*const N as &'a N));

slice_storage_impl!("A mutable matrix data storage for mutable matrix slice. Only contains an \
                     internal mutable reference to another matrix data storage.";
    StorageMut as &'a mut S; SliceStorageMut.get_address_unchecked_mut(*mut N as &'a mut N)
);


impl<'a, N: Scalar, R: Dim, C: Dim, RStride: Dim, CStride: Dim, Alloc> Copy
for SliceStorage<'a, N, R, C, RStride, CStride, Alloc> { }

impl<'a, N: Scalar, R: Dim, C: Dim, RStride: Dim, CStride: Dim, Alloc> Clone
for SliceStorage<'a, N, R, C, RStride, CStride, Alloc> {
    #[inline]
    fn clone(&self) -> Self {
        SliceStorage {
            ptr:       self.ptr,
            shape:     self.shape,
            strides:   self.strides,
            _phantoms: PhantomData,
        }
    }
}

macro_rules! storage_impl(
    ($($T: ident),* $(,)*) => {$(
        unsafe impl<'a, N, R: Dim, C: Dim, RStride: Dim, CStride: Dim, Alloc> Storage<N, R, C>
            for $T<'a, N, R, C, RStride, CStride, Alloc>
            where N:     Scalar,
                  Alloc: Allocator<N, R, C> {

            type RStride = RStride;
            type CStride = CStride;
            type Alloc   = Alloc;

            #[inline]
            fn into_owned(self) -> Owned<N, R, C, Self::Alloc> {
                self.clone_owned()
            }

            #[inline]
            fn clone_owned(&self) -> Owned<N, R, C, Self::Alloc> {
                let (nrows, ncols) = self.shape();
                let it = MatrixIter::new(self).cloned();
                Alloc::allocate_from_iterator(nrows, ncols, it)
            }

            #[inline]
            fn ptr(&self) -> *const N {
                self.ptr
            }

            #[inline]
            fn shape(&self) -> (R, C) {
                self.shape
            }

            #[inline]
            fn strides(&self) -> (Self::RStride, Self::CStride) {
                self.strides
            }
        }
    )*}
);

storage_impl!(SliceStorage, SliceStorageMut);

unsafe impl<'a, N, R: Dim, C: Dim, RStride: Dim, CStride: Dim, Alloc> StorageMut<N, R, C>
    for SliceStorageMut<'a, N, R, C, RStride, CStride, Alloc>
    where N: Scalar,
          Alloc: Allocator<N, R, C> {
    #[inline]
    fn ptr_mut(&mut self) -> *mut N {
        self.ptr
    }
}


impl<N: Scalar, R: Dim, C: Dim, S: Storage<N, R, C>> Matrix<N, R, C, S> {
    #[inline]
    fn assert_slice_index(&self, start: (usize, usize), shape: (usize, usize), steps: (usize, usize)) {
        let my_shape = self.shape();
        assert!(start.0 + (shape.0 - 1) * steps.0 <= my_shape.0, "Matrix slicing out of bounds.");
        assert!(start.1 + (shape.1 - 1) * steps.1 <= my_shape.1, "Matrix slicing out of bounds.");
    }
}


macro_rules! matrix_slice_impl(
    ($me: ident: $Me: ty, $MatrixSlice: ident, $SliceStorage: ident, $Storage: ident, $data: expr;
     $row: ident,
     $rows: ident,
     $rows_with_step: ident,
     $fixed_rows: ident,
     $fixed_rows_with_step: ident,
     $rows_generic: ident,
     $column: ident,
     $columns: ident,
     $columns_with_step: ident,
     $fixed_columns: ident,
     $fixed_columns_with_step: ident,
     $columns_generic: ident,
     $slice: ident,
     $slice_with_steps: ident,
     $fixed_slice: ident,
     $fixed_slice_with_steps: ident,
     $generic_slice: ident,
     $generic_slice_with_steps: ident) => {
        /// A matrix slice.
        pub type $MatrixSlice<'a, N, R, C, RStride, CStride, Alloc>
        = Matrix<N, R, C, $SliceStorage<'a, N, R, C, RStride, CStride, Alloc>>;

        impl<N: Scalar, R: Dim, C: Dim, S: $Storage<N, R, C>> Matrix<N, R, C, S> {
            /*
             *
             * Row slicing.
             *
             */
            /// Returns a slice containing the i-th column of this matrix.
            #[inline]
            pub fn $row($me: $Me, i: usize) -> $MatrixSlice<N, U1, C, S::RStride, S::CStride, S::Alloc> {
                $me.$fixed_rows::<U1>(i)
            }

            /// Extracts from this matrix a set of consecutive rows.
            #[inline]
            pub fn $rows($me: $Me, first_row: usize, nrows: usize)
                -> $MatrixSlice<N, Dynamic, C, S::RStride, S::CStride, S::Alloc> {

                let my_shape = $me.data.shape();
                $me.assert_slice_index((first_row, 0), (nrows, my_shape.1.value()), (1, 1));
                let shape = (Dynamic::new(nrows), my_shape.1);

                unsafe {
                    let data = $SliceStorage::new_unchecked($data, (first_row, 0), shape);
                    Matrix::from_data_statically_unchecked(data)
                }
            }

            /// Extracts from this matrix a set of consecutive rows regularly spaced by `step` rows.
            #[inline]
            pub fn $rows_with_step($me: $Me, first_row: usize, nrows: usize, step: usize)
                -> $MatrixSlice<N, Dynamic, C, Dynamic, S::CStride, S::Alloc> {

                $me.$rows_generic(first_row, Dynamic::new(nrows), Dynamic::new(step))
            }

            /// Extracts a compile-time number of consecutive rows from this matrix.
            #[inline]
            pub fn $fixed_rows<RSlice>($me: $Me, first_row: usize)
                -> $MatrixSlice<N, RSlice, C, S::RStride, S::CStride, S::Alloc>
                where RSlice: DimName {

                let my_shape = $me.data.shape();
                $me.assert_slice_index((first_row, 0), (RSlice::dim(), my_shape.1.value()), (1, 1));
                let shape = (RSlice::name(), my_shape.1);

                unsafe {
                    let data = $SliceStorage::new_unchecked($data, (first_row, 0), shape);
                    Matrix::from_data_statically_unchecked(data)
                }
            }

            /// Extracts from this matrix a compile-time number of rows regularly spaced by `step` rows.
            #[inline]
            pub fn $fixed_rows_with_step<RSlice>($me: $Me, first_row: usize, step: usize)
                -> $MatrixSlice<N, RSlice, C, Dynamic, S::CStride, S::Alloc>
                where RSlice: DimName {

                $me.$rows_generic(first_row, RSlice::name(), Dynamic::new(step))
            }

            /// Extracts from this matrix `nrows` rows regularly spaced by `step` rows. Both argument may
            /// or may not be values known at compile-time.
            #[inline]
            pub fn $rows_generic<RSlice, RStep>($me: $Me, row_start: usize, nrows: RSlice, step: RStep)
                -> $MatrixSlice<N, RSlice, C, DimProd<RStep, S::RStride>, S::CStride, S::Alloc>
                where RSlice: Dim,
                      RStep: DimMul<S::RStride> {

                let my_shape   = $me.data.shape();
                let my_strides = $me.data.strides();
                $me.assert_slice_index((row_start, 0), (nrows.value(), my_shape.1.value()), (step.value(), 1));

                let strides = (step.mul(my_strides.0), my_strides.1);
                let shape   = (nrows, my_shape.1);

                unsafe {
                    let data = $SliceStorage::new_with_strides_unchecked($data, (row_start, 0), shape, strides);
                    Matrix::from_data_statically_unchecked(data)
                }
            }

            /*
             *
             * Column slicing.
             *
             */
            /// Returns a slice containing the i-th column of this matrix.
            #[inline]
            pub fn $column($me: $Me, i: usize) -> $MatrixSlice<N, R, U1, S::RStride, S::CStride, S::Alloc> {
                $me.$fixed_columns::<U1>(i)
            }

            /// Extracts from this matrix a set of consecutive columns.
            #[inline]
            pub fn $columns($me: $Me, first_col: usize, ncols: usize)
                -> $MatrixSlice<N, R, Dynamic, S::RStride, S::CStride, S::Alloc> {

                let my_shape = $me.data.shape();
                $me.assert_slice_index((0, first_col), (my_shape.0.value(), ncols), (1, 1));
                let shape = (my_shape.0, Dynamic::new(ncols));

                unsafe {
                    let data = $SliceStorage::new_unchecked($data, (0, first_col), shape);
                    Matrix::from_data_statically_unchecked(data)
                }
            }

            /// Extracts from this matrix a set of consecutive columns regularly spaced by `step` columns.
            #[inline]
            pub fn $columns_with_step($me: $Me, first_col: usize, ncols: usize, step: usize)
                -> $MatrixSlice<N, R, Dynamic, S::RStride, Dynamic, S::Alloc> {

                $me.$columns_generic(first_col, Dynamic::new(ncols), Dynamic::new(step))
            }

            /// Extracts a compile-time number of consecutive columns from this matrix.
            #[inline]
            pub fn $fixed_columns<CSlice>($me: $Me, first_col: usize)
                -> $MatrixSlice<N, R, CSlice, S::RStride, S::CStride, S::Alloc>
                where CSlice: DimName {

                let my_shape = $me.data.shape();
                $me.assert_slice_index((0, first_col), (my_shape.0.value(), CSlice::dim()), (1, 1));
                let shape = (my_shape.0, CSlice::name());

                unsafe {
                    let data = $SliceStorage::new_unchecked($data, (0, first_col), shape);
                    Matrix::from_data_statically_unchecked(data)
                }
            }

            /// Extracts from this matrix a compile-time number of columns regularly spaced by `step`
            /// columns.
            #[inline]
            pub fn $fixed_columns_with_step<CSlice>($me: $Me, first_col: usize, step: usize)
                -> $MatrixSlice<N, R, CSlice, S::RStride, Dynamic, S::Alloc>
                where CSlice: DimName {

                $me.$columns_generic(first_col, CSlice::name(), Dynamic::new(step))
            }

            /// Extracts from this matrix `ncols` columns regularly spaced by `step` columns. Both argument may
            /// or may not be values known at compile-time.
            #[inline]
            pub fn $columns_generic<CSlice, CStep>($me: $Me, first_col: usize, ncols: CSlice, step: CStep)
                -> $MatrixSlice<N, R, CSlice, S::RStride, DimProd<CStep, S::CStride>, S::Alloc>
                where CSlice: Dim,
                      CStep: DimMul<S::CStride> {

                let my_shape   = $me.data.shape();
                let my_strides = $me.data.strides();

                $me.assert_slice_index((0, first_col), (my_shape.0.value(), ncols.value()), (1, step.value()));

                let strides = (my_strides.0, step.mul(my_strides.1));
                let shape   = (my_shape.0, ncols);

                unsafe {
                    let data = $SliceStorage::new_with_strides_unchecked($data, (0, first_col), shape, strides);
                    Matrix::from_data_statically_unchecked(data)
                }
            }

            /*
             *
             * General slicing.
             *
             */
            /// Slices this matrix starting at its component `(irow, icol)` and with `(nrows, ncols)`
            /// consecutive elements.
            #[inline]
            pub fn $slice($me: $Me, start: (usize, usize), shape: (usize, usize))
                -> $MatrixSlice<N, Dynamic, Dynamic, S::RStride, S::CStride, S::Alloc> {

                $me.assert_slice_index(start, shape, (1, 1));
                let shape = (Dynamic::new(shape.0), Dynamic::new(shape.1));

                unsafe {
                    let data = $SliceStorage::new_unchecked($data, start, shape);
                    Matrix::from_data_statically_unchecked(data)
                }
            }


            /// Slices this matrix starting at its component `(start.0, start.1)` and with
            /// `(shape.0, shape.1)` components. Each row (resp. column) of the sliced matrix is
            /// separated by `steps.0` (resp. `steps.1`) ignored rows (resp. columns) of the
            /// original matrix.
            #[inline]
            pub fn $slice_with_steps($me: $Me, start: (usize, usize), shape: (usize, usize), steps: (usize, usize))
                -> $MatrixSlice<N, Dynamic, Dynamic, Dynamic, Dynamic, S::Alloc> {
                let shape = (Dynamic::new(shape.0), Dynamic::new(shape.1));
                let steps = (Dynamic::new(steps.0), Dynamic::new(steps.1));

                $me.$generic_slice_with_steps(start, shape, steps)
            }

            /// Slices this matrix starting at its component `(irow, icol)` and with `(R::dim(),
            /// CSlice::dim())` consecutive components.
            #[inline]
            pub fn $fixed_slice<RSlice, CSlice>($me: $Me, irow: usize, icol: usize)
                -> $MatrixSlice<N, RSlice, CSlice, S::RStride, S::CStride, S::Alloc>
                where RSlice: DimName,
                      CSlice: DimName {

                $me.assert_slice_index((irow, icol), (RSlice::dim(), CSlice::dim()), (1, 1));
                let shape = (RSlice::name(), CSlice::name());

                unsafe {
                    let data = $SliceStorage::new_unchecked($data, (irow, icol), shape);
                    Matrix::from_data_statically_unchecked(data)
                }
            }

            /// Slices this matrix starting at its component `(start.0, start.1)` and with
            /// `(R::dim(), CSlice::dim())` components. Each row (resp. column) of the sliced
            /// matrix is separated by `steps.0` (resp. `steps.1`) ignored rows (resp. columns) of
            /// the original matrix.
            #[inline]
            pub fn $fixed_slice_with_steps<RSlice, CSlice>($me: $Me, start: (usize, usize), steps: (usize, usize))
                -> $MatrixSlice<N, RSlice, CSlice, Dynamic, Dynamic, S::Alloc>
                where RSlice: DimName,
                      CSlice: DimName {
                let shape = (RSlice::name(), CSlice::name());
                let steps = (Dynamic::new(steps.0), Dynamic::new(steps.1));
                $me.$generic_slice_with_steps(start, shape, steps)
            }

            /// Creates a slice that may or may not have a fixed size and stride.
            #[inline]
            pub fn $generic_slice<RSlice, CSlice>($me: $Me, start: (usize, usize), shape: (RSlice, CSlice))
                -> $MatrixSlice<N, RSlice, CSlice, S::RStride, S::CStride, S::Alloc>
                where RSlice: Dim,
                      CSlice: Dim {

                $me.assert_slice_index(start, (shape.0.value(), shape.1.value()), (1, 1));

                unsafe {
                    let data = $SliceStorage::new_unchecked($data, start, shape);
                    Matrix::from_data_statically_unchecked(data)
                }
            }

            /// Creates a slice that may or may not have a fixed size and stride.
            #[inline]
            pub fn $generic_slice_with_steps<RSlice, CSlice, RStep, CStep>($me: $Me,
                                                                           start: (usize, usize),
                                                                           shape: (RSlice, CSlice),
                                                                           steps: (RStep, CStep))
                -> $MatrixSlice<N, RSlice, CSlice, DimProd<RStep, S::RStride>, DimProd<CStep, S::CStride>, S::Alloc>
                where RSlice: Dim,
                      CSlice: Dim,
                      RStep: DimMul<S::RStride>,
                      CStep: DimMul<S::CStride> {

                $me.assert_slice_index(start, (shape.0.value(), shape.1.value()), (steps.0.value(), steps.1.value()));

                let my_strides = $me.data.strides();
                let strides    = (steps.0.mul(my_strides.0), steps.1.mul(my_strides.1));

                unsafe {
                    let data = $SliceStorage::new_with_strides_unchecked($data, start, shape, strides);
                    Matrix::from_data_statically_unchecked(data)
                }
            }
        }
    }
);

matrix_slice_impl!(
     self: &Self, MatrixSlice, SliceStorage, Storage, &self.data;
     row,
     rows,
     rows_with_step,
     fixed_rows,
     fixed_rows_with_step,
     rows_generic,
     column,
     columns,
     columns_with_step,
     fixed_columns,
     fixed_columns_with_step,
     columns_generic,
     slice,
     slice_with_steps,
     fixed_slice,
     fixed_slice_with_steps,
     generic_slice,
     generic_slice_with_steps);


matrix_slice_impl!(
     self: &mut Self, MatrixSliceMut, SliceStorageMut, StorageMut, &mut self.data;
     row_mut,
     rows_mut,
     rows_with_step_mut,
     fixed_rows_mut,
     fixed_rows_with_step_mut,
     rows_generic_mut,
     column_mut,
     columns_mut,
     columns_with_step_mut,
     fixed_columns_mut,
     fixed_columns_with_step_mut,
     columns_generic_mut,
     slice_mut,
     slice_with_steps_mut,
     fixed_slice_mut,
     fixed_slice_with_steps_mut,
     generic_slice_mut,
     generic_slice_with_steps_mut);