use crate::dimension::slices_intersect;
use crate::error::{ErrorKind, ShapeError};
use crate::{ArrayViewMut, DimAdd, Dimension, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn};
use alloc::vec::Vec;
use std::convert::TryFrom;
use std::fmt;
use std::marker::PhantomData;
use std::ops::{Deref, Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive};
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub struct Slice {
pub start: isize,
pub end: Option<isize>,
pub step: isize,
}
impl Slice {
pub fn new(start: isize, end: Option<isize>, step: isize) -> Slice {
debug_assert_ne!(step, 0, "Slice::new: step must be nonzero");
Slice { start, end, step }
}
#[inline]
pub fn step_by(self, step: isize) -> Self {
debug_assert_ne!(step, 0, "Slice::step_by: step must be nonzero");
Slice {
step: self.step * step,
..self
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct NewAxis;
#[derive(Debug, PartialEq, Eq, Hash)]
pub enum SliceInfoElem {
Slice {
start: isize,
end: Option<isize>,
step: isize,
},
Index(isize),
NewAxis,
}
copy_and_clone! {SliceInfoElem}
impl SliceInfoElem {
pub fn is_slice(&self) -> bool {
matches!(self, SliceInfoElem::Slice { .. })
}
pub fn is_index(&self) -> bool {
matches!(self, SliceInfoElem::Index(_))
}
pub fn is_new_axis(&self) -> bool {
matches!(self, SliceInfoElem::NewAxis)
}
}
impl fmt::Display for SliceInfoElem {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match *self {
SliceInfoElem::Index(index) => write!(f, "{}", index)?,
SliceInfoElem::Slice { start, end, step } => {
if start != 0 {
write!(f, "{}", start)?;
}
write!(f, "..")?;
if let Some(i) = end {
write!(f, "{}", i)?;
}
if step != 1 {
write!(f, ";{}", step)?;
}
}
SliceInfoElem::NewAxis => write!(f, stringify!(NewAxis))?,
}
Ok(())
}
}
macro_rules! impl_slice_variant_from_range {
($self:ty, $constructor:path, $index:ty) => {
impl From<Range<$index>> for $self {
#[inline]
fn from(r: Range<$index>) -> $self {
$constructor {
start: r.start as isize,
end: Some(r.end as isize),
step: 1,
}
}
}
impl From<RangeInclusive<$index>> for $self {
#[inline]
fn from(r: RangeInclusive<$index>) -> $self {
let end = *r.end() as isize;
$constructor {
start: *r.start() as isize,
end: if end == -1 { None } else { Some(end + 1) },
step: 1,
}
}
}
impl From<RangeFrom<$index>> for $self {
#[inline]
fn from(r: RangeFrom<$index>) -> $self {
$constructor {
start: r.start as isize,
end: None,
step: 1,
}
}
}
impl From<RangeTo<$index>> for $self {
#[inline]
fn from(r: RangeTo<$index>) -> $self {
$constructor {
start: 0,
end: Some(r.end as isize),
step: 1,
}
}
}
impl From<RangeToInclusive<$index>> for $self {
#[inline]
fn from(r: RangeToInclusive<$index>) -> $self {
let end = r.end as isize;
$constructor {
start: 0,
end: if end == -1 { None } else { Some(end + 1) },
step: 1,
}
}
}
};
}
impl_slice_variant_from_range!(Slice, Slice, isize);
impl_slice_variant_from_range!(Slice, Slice, usize);
impl_slice_variant_from_range!(Slice, Slice, i32);
impl_slice_variant_from_range!(SliceInfoElem, SliceInfoElem::Slice, isize);
impl_slice_variant_from_range!(SliceInfoElem, SliceInfoElem::Slice, usize);
impl_slice_variant_from_range!(SliceInfoElem, SliceInfoElem::Slice, i32);
impl From<RangeFull> for Slice {
#[inline]
fn from(_: RangeFull) -> Slice {
Slice {
start: 0,
end: None,
step: 1,
}
}
}
impl From<RangeFull> for SliceInfoElem {
#[inline]
fn from(_: RangeFull) -> SliceInfoElem {
SliceInfoElem::Slice {
start: 0,
end: None,
step: 1,
}
}
}
impl From<Slice> for SliceInfoElem {
#[inline]
fn from(s: Slice) -> SliceInfoElem {
SliceInfoElem::Slice {
start: s.start,
end: s.end,
step: s.step,
}
}
}
macro_rules! impl_sliceinfoelem_from_index {
($index:ty) => {
impl From<$index> for SliceInfoElem {
#[inline]
fn from(r: $index) -> SliceInfoElem {
SliceInfoElem::Index(r as isize)
}
}
};
}
impl_sliceinfoelem_from_index!(isize);
impl_sliceinfoelem_from_index!(usize);
impl_sliceinfoelem_from_index!(i32);
impl From<NewAxis> for SliceInfoElem {
#[inline]
fn from(_: NewAxis) -> SliceInfoElem {
SliceInfoElem::NewAxis
}
}
pub unsafe trait SliceArg<D: Dimension>: AsRef<[SliceInfoElem]> {
type OutDim: Dimension;
fn in_ndim(&self) -> usize;
fn out_ndim(&self) -> usize;
private_decl! {}
}
unsafe impl<T, D> SliceArg<D> for &T
where
T: SliceArg<D> + ?Sized,
D: Dimension,
{
type OutDim = T::OutDim;
fn in_ndim(&self) -> usize {
T::in_ndim(self)
}
fn out_ndim(&self) -> usize {
T::out_ndim(self)
}
private_impl! {}
}
macro_rules! impl_slicearg_samedim {
($in_dim:ty) => {
unsafe impl<T, Dout> SliceArg<$in_dim> for SliceInfo<T, $in_dim, Dout>
where
T: AsRef<[SliceInfoElem]>,
Dout: Dimension,
{
type OutDim = Dout;
fn in_ndim(&self) -> usize {
self.in_ndim()
}
fn out_ndim(&self) -> usize {
self.out_ndim()
}
private_impl! {}
}
};
}
impl_slicearg_samedim!(Ix0);
impl_slicearg_samedim!(Ix1);
impl_slicearg_samedim!(Ix2);
impl_slicearg_samedim!(Ix3);
impl_slicearg_samedim!(Ix4);
impl_slicearg_samedim!(Ix5);
impl_slicearg_samedim!(Ix6);
unsafe impl<T, Din, Dout> SliceArg<IxDyn> for SliceInfo<T, Din, Dout>
where
T: AsRef<[SliceInfoElem]>,
Din: Dimension,
Dout: Dimension,
{
type OutDim = Dout;
fn in_ndim(&self) -> usize {
self.in_ndim()
}
fn out_ndim(&self) -> usize {
self.out_ndim()
}
private_impl! {}
}
unsafe impl SliceArg<IxDyn> for [SliceInfoElem] {
type OutDim = IxDyn;
fn in_ndim(&self) -> usize {
self.iter().filter(|s| !s.is_new_axis()).count()
}
fn out_ndim(&self) -> usize {
self.iter().filter(|s| !s.is_index()).count()
}
private_impl! {}
}
#[derive(Debug)]
pub struct SliceInfo<T, Din: Dimension, Dout: Dimension> {
in_dim: PhantomData<Din>,
out_dim: PhantomData<Dout>,
indices: T,
}
impl<T, Din, Dout> Deref for SliceInfo<T, Din, Dout>
where
Din: Dimension,
Dout: Dimension,
{
type Target = T;
fn deref(&self) -> &Self::Target {
&self.indices
}
}
fn check_dims_for_sliceinfo<Din, Dout>(indices: &[SliceInfoElem]) -> Result<(), ShapeError>
where
Din: Dimension,
Dout: Dimension,
{
if let Some(in_ndim) = Din::NDIM {
if in_ndim != indices.in_ndim() {
return Err(ShapeError::from_kind(ErrorKind::IncompatibleShape));
}
}
if let Some(out_ndim) = Dout::NDIM {
if out_ndim != indices.out_ndim() {
return Err(ShapeError::from_kind(ErrorKind::IncompatibleShape));
}
}
Ok(())
}
impl<T, Din, Dout> SliceInfo<T, Din, Dout>
where
T: AsRef<[SliceInfoElem]>,
Din: Dimension,
Dout: Dimension,
{
#[doc(hidden)]
pub unsafe fn new_unchecked(
indices: T,
in_dim: PhantomData<Din>,
out_dim: PhantomData<Dout>,
) -> SliceInfo<T, Din, Dout> {
if cfg!(debug_assertions) {
check_dims_for_sliceinfo::<Din, Dout>(indices.as_ref())
.expect("`Din` and `Dout` must be consistent with `indices`.");
}
SliceInfo {
in_dim,
out_dim,
indices,
}
}
pub unsafe fn new(indices: T) -> Result<SliceInfo<T, Din, Dout>, ShapeError> {
check_dims_for_sliceinfo::<Din, Dout>(indices.as_ref())?;
Ok(SliceInfo {
in_dim: PhantomData,
out_dim: PhantomData,
indices,
})
}
pub fn in_ndim(&self) -> usize {
if let Some(ndim) = Din::NDIM {
ndim
} else {
self.indices.as_ref().in_ndim()
}
}
pub fn out_ndim(&self) -> usize {
if let Some(ndim) = Dout::NDIM {
ndim
} else {
self.indices.as_ref().out_ndim()
}
}
}
impl<'a, Din, Dout> TryFrom<&'a [SliceInfoElem]> for SliceInfo<&'a [SliceInfoElem], Din, Dout>
where
Din: Dimension,
Dout: Dimension,
{
type Error = ShapeError;
fn try_from(
indices: &'a [SliceInfoElem],
) -> Result<SliceInfo<&'a [SliceInfoElem], Din, Dout>, ShapeError> {
unsafe {
Self::new(indices)
}
}
}
impl<Din, Dout> TryFrom<Vec<SliceInfoElem>> for SliceInfo<Vec<SliceInfoElem>, Din, Dout>
where
Din: Dimension,
Dout: Dimension,
{
type Error = ShapeError;
fn try_from(
indices: Vec<SliceInfoElem>,
) -> Result<SliceInfo<Vec<SliceInfoElem>, Din, Dout>, ShapeError> {
unsafe {
Self::new(indices)
}
}
}
macro_rules! impl_tryfrom_array_for_sliceinfo {
($len:expr) => {
impl<Din, Dout> TryFrom<[SliceInfoElem; $len]>
for SliceInfo<[SliceInfoElem; $len], Din, Dout>
where
Din: Dimension,
Dout: Dimension,
{
type Error = ShapeError;
fn try_from(
indices: [SliceInfoElem; $len],
) -> Result<SliceInfo<[SliceInfoElem; $len], Din, Dout>, ShapeError> {
unsafe {
Self::new(indices)
}
}
}
};
}
impl_tryfrom_array_for_sliceinfo!(0);
impl_tryfrom_array_for_sliceinfo!(1);
impl_tryfrom_array_for_sliceinfo!(2);
impl_tryfrom_array_for_sliceinfo!(3);
impl_tryfrom_array_for_sliceinfo!(4);
impl_tryfrom_array_for_sliceinfo!(5);
impl_tryfrom_array_for_sliceinfo!(6);
impl_tryfrom_array_for_sliceinfo!(7);
impl_tryfrom_array_for_sliceinfo!(8);
impl<T, Din, Dout> AsRef<[SliceInfoElem]> for SliceInfo<T, Din, Dout>
where
T: AsRef<[SliceInfoElem]>,
Din: Dimension,
Dout: Dimension,
{
fn as_ref(&self) -> &[SliceInfoElem] {
self.indices.as_ref()
}
}
impl<'a, T, Din, Dout> From<&'a SliceInfo<T, Din, Dout>>
for SliceInfo<&'a [SliceInfoElem], Din, Dout>
where
T: AsRef<[SliceInfoElem]>,
Din: Dimension,
Dout: Dimension,
{
fn from(info: &'a SliceInfo<T, Din, Dout>) -> SliceInfo<&'a [SliceInfoElem], Din, Dout> {
SliceInfo {
in_dim: info.in_dim,
out_dim: info.out_dim,
indices: info.indices.as_ref(),
}
}
}
impl<T, Din, Dout> Copy for SliceInfo<T, Din, Dout>
where
T: Copy,
Din: Dimension,
Dout: Dimension,
{
}
impl<T, Din, Dout> Clone for SliceInfo<T, Din, Dout>
where
T: Clone,
Din: Dimension,
Dout: Dimension,
{
fn clone(&self) -> Self {
SliceInfo {
in_dim: PhantomData,
out_dim: PhantomData,
indices: self.indices.clone(),
}
}
}
#[doc(hidden)]
pub trait SliceNextDim {
type InDim: Dimension;
type OutDim: Dimension;
fn next_in_dim<D>(&self, _: PhantomData<D>) -> PhantomData<<D as DimAdd<Self::InDim>>::Output>
where
D: Dimension + DimAdd<Self::InDim>,
{
PhantomData
}
fn next_out_dim<D>(&self, _: PhantomData<D>) -> PhantomData<<D as DimAdd<Self::OutDim>>::Output>
where
D: Dimension + DimAdd<Self::OutDim>,
{
PhantomData
}
}
macro_rules! impl_slicenextdim {
(($($generics:tt)*), $self:ty, $in:ty, $out:ty) => {
impl<$($generics)*> SliceNextDim for $self {
type InDim = $in;
type OutDim = $out;
}
};
}
impl_slicenextdim!((), isize, Ix1, Ix0);
impl_slicenextdim!((), usize, Ix1, Ix0);
impl_slicenextdim!((), i32, Ix1, Ix0);
impl_slicenextdim!((T), Range<T>, Ix1, Ix1);
impl_slicenextdim!((T), RangeInclusive<T>, Ix1, Ix1);
impl_slicenextdim!((T), RangeFrom<T>, Ix1, Ix1);
impl_slicenextdim!((T), RangeTo<T>, Ix1, Ix1);
impl_slicenextdim!((T), RangeToInclusive<T>, Ix1, Ix1);
impl_slicenextdim!((), RangeFull, Ix1, Ix1);
impl_slicenextdim!((), Slice, Ix1, Ix1);
impl_slicenextdim!((), NewAxis, Ix0, Ix1);
#[macro_export]
macro_rules! s(
(@parse $in_dim:expr, $out_dim:expr, [$($stack:tt)*] $r:expr;$s:expr) => {
match $r {
r => {
let in_dim = $crate::SliceNextDim::next_in_dim(&r, $in_dim);
let out_dim = $crate::SliceNextDim::next_out_dim(&r, $out_dim);
#[allow(unsafe_code)]
unsafe {
$crate::SliceInfo::new_unchecked(
[$($stack)* $crate::s!(@convert r, $s)],
in_dim,
out_dim,
)
}
}
}
};
(@parse $in_dim:expr, $out_dim:expr, [$($stack:tt)*] $r:expr) => {
match $r {
r => {
let in_dim = $crate::SliceNextDim::next_in_dim(&r, $in_dim);
let out_dim = $crate::SliceNextDim::next_out_dim(&r, $out_dim);
#[allow(unsafe_code)]
unsafe {
$crate::SliceInfo::new_unchecked(
[$($stack)* $crate::s!(@convert r)],
in_dim,
out_dim,
)
}
}
}
};
(@parse $in_dim:expr, $out_dim:expr, [$($stack:tt)*] $r:expr;$s:expr ,) => {
$crate::s![@parse $in_dim, $out_dim, [$($stack)*] $r;$s]
};
(@parse $in_dim:expr, $out_dim:expr, [$($stack:tt)*] $r:expr ,) => {
$crate::s![@parse $in_dim, $out_dim, [$($stack)*] $r]
};
(@parse $in_dim:expr, $out_dim:expr, [$($stack:tt)*] $r:expr;$s:expr, $($t:tt)*) => {
match $r {
r => {
$crate::s![@parse
$crate::SliceNextDim::next_in_dim(&r, $in_dim),
$crate::SliceNextDim::next_out_dim(&r, $out_dim),
[$($stack)* $crate::s!(@convert r, $s),]
$($t)*
]
}
}
};
(@parse $in_dim:expr, $out_dim:expr, [$($stack:tt)*] $r:expr, $($t:tt)*) => {
match $r {
r => {
$crate::s![@parse
$crate::SliceNextDim::next_in_dim(&r, $in_dim),
$crate::SliceNextDim::next_out_dim(&r, $out_dim),
[$($stack)* $crate::s!(@convert r),]
$($t)*
]
}
}
};
(@parse ::std::marker::PhantomData::<$crate::Ix0>, ::std::marker::PhantomData::<$crate::Ix0>, []) => {
{
#[allow(unsafe_code)]
unsafe {
$crate::SliceInfo::new_unchecked(
[],
::std::marker::PhantomData::<$crate::Ix0>,
::std::marker::PhantomData::<$crate::Ix0>,
)
}
}
};
(@parse $($t:tt)*) => { compile_error!("Invalid syntax in s![] call.") };
(@convert $r:expr) => {
<$crate::SliceInfoElem as ::std::convert::From<_>>::from($r)
};
(@convert $r:expr, $s:expr) => {
<$crate::SliceInfoElem as ::std::convert::From<_>>::from(
<$crate::Slice as ::std::convert::From<_>>::from($r).step_by($s as isize)
)
};
($($t:tt)*) => {
$crate::s![@parse
::std::marker::PhantomData::<$crate::Ix0>,
::std::marker::PhantomData::<$crate::Ix0>,
[]
$($t)*
]
};
);
pub trait MultiSliceArg<'a, A, D>
where
A: 'a,
D: Dimension,
{
type Output;
fn multi_slice_move(&self, view: ArrayViewMut<'a, A, D>) -> Self::Output;
private_decl! {}
}
impl<'a, A, D> MultiSliceArg<'a, A, D> for ()
where
A: 'a,
D: Dimension,
{
type Output = ();
fn multi_slice_move(&self, _view: ArrayViewMut<'a, A, D>) -> Self::Output {}
private_impl! {}
}
impl<'a, A, D, I0> MultiSliceArg<'a, A, D> for (I0,)
where
A: 'a,
D: Dimension,
I0: SliceArg<D>,
{
type Output = (ArrayViewMut<'a, A, I0::OutDim>,);
fn multi_slice_move(&self, view: ArrayViewMut<'a, A, D>) -> Self::Output {
(view.slice_move(&self.0),)
}
private_impl! {}
}
macro_rules! impl_multislice_tuple {
([$($but_last:ident)*] $last:ident) => {
impl_multislice_tuple!(@def_impl ($($but_last,)* $last,), [$($but_last)*] $last);
};
(@def_impl ($($all:ident,)*), [$($but_last:ident)*] $last:ident) => {
impl<'a, A, D, $($all,)*> MultiSliceArg<'a, A, D> for ($($all,)*)
where
A: 'a,
D: Dimension,
$($all: SliceArg<D>,)*
{
type Output = ($(ArrayViewMut<'a, A, $all::OutDim>,)*);
fn multi_slice_move(&self, view: ArrayViewMut<'a, A, D>) -> Self::Output {
#[allow(non_snake_case)]
let ($($all,)*) = self;
let shape = view.raw_dim();
assert!(!impl_multislice_tuple!(@intersects_self &shape, ($($all,)*)));
let raw_view = view.into_raw_view_mut();
unsafe {
(
$(raw_view.clone().slice_move($but_last).deref_into_view_mut(),)*
raw_view.slice_move($last).deref_into_view_mut(),
)
}
}
private_impl! {}
}
};
(@intersects_self $shape:expr, ($head:expr,)) => {
false
};
(@intersects_self $shape:expr, ($head:expr, $($tail:expr,)*)) => {
$(slices_intersect($shape, $head, $tail)) ||*
|| impl_multislice_tuple!(@intersects_self $shape, ($($tail,)*))
};
}
impl_multislice_tuple!([I0] I1);
impl_multislice_tuple!([I0 I1] I2);
impl_multislice_tuple!([I0 I1 I2] I3);
impl_multislice_tuple!([I0 I1 I2 I3] I4);
impl_multislice_tuple!([I0 I1 I2 I3 I4] I5);
impl<'a, A, D, T> MultiSliceArg<'a, A, D> for &T
where
A: 'a,
D: Dimension,
T: MultiSliceArg<'a, A, D>,
{
type Output = T::Output;
fn multi_slice_move(&self, view: ArrayViewMut<'a, A, D>) -> Self::Output {
T::multi_slice_move(self, view)
}
private_impl! {}
}