use crate::dimension::slices_intersect;
use crate::error::{ErrorKind, ShapeError};
use crate::{ArrayViewMut, Dimension};
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(Debug, PartialEq, Eq, Hash)]
pub enum SliceOrIndex {
Slice {
start: isize,
end: Option<isize>,
step: isize,
},
Index(isize),
}
copy_and_clone! {SliceOrIndex}
impl SliceOrIndex {
pub fn is_slice(&self) -> bool {
match self {
SliceOrIndex::Slice { .. } => true,
_ => false,
}
}
pub fn is_index(&self) -> bool {
match self {
SliceOrIndex::Index(_) => true,
_ => false,
}
}
#[inline]
pub fn step_by(self, step: isize) -> Self {
debug_assert_ne!(step, 0, "SliceOrIndex::step_by: step must be nonzero");
match self {
SliceOrIndex::Slice {
start,
end,
step: orig_step,
} => SliceOrIndex::Slice {
start,
end,
step: orig_step * step,
},
SliceOrIndex::Index(s) => SliceOrIndex::Index(s),
}
}
}
impl fmt::Display for SliceOrIndex {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match *self {
SliceOrIndex::Index(index) => write!(f, "{}", index)?,
SliceOrIndex::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)?;
}
}
}
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!(SliceOrIndex, SliceOrIndex::Slice, isize);
impl_slice_variant_from_range!(SliceOrIndex, SliceOrIndex::Slice, usize);
impl_slice_variant_from_range!(SliceOrIndex, SliceOrIndex::Slice, i32);
impl From<RangeFull> for Slice {
#[inline]
fn from(_: RangeFull) -> Slice {
Slice {
start: 0,
end: None,
step: 1,
}
}
}
impl From<RangeFull> for SliceOrIndex {
#[inline]
fn from(_: RangeFull) -> SliceOrIndex {
SliceOrIndex::Slice {
start: 0,
end: None,
step: 1,
}
}
}
impl From<Slice> for SliceOrIndex {
#[inline]
fn from(s: Slice) -> SliceOrIndex {
SliceOrIndex::Slice {
start: s.start,
end: s.end,
step: s.step,
}
}
}
macro_rules! impl_sliceorindex_from_index {
($index:ty) => {
impl From<$index> for SliceOrIndex {
#[inline]
fn from(r: $index) -> SliceOrIndex {
SliceOrIndex::Index(r as isize)
}
}
};
}
impl_sliceorindex_from_index!(isize);
impl_sliceorindex_from_index!(usize);
impl_sliceorindex_from_index!(i32);
#[derive(Debug)]
#[repr(C)]
pub struct SliceInfo<T: ?Sized, D: Dimension> {
out_dim: PhantomData<D>,
indices: T,
}
impl<T: ?Sized, D> Deref for SliceInfo<T, D>
where
D: Dimension,
{
type Target = T;
fn deref(&self) -> &Self::Target {
&self.indices
}
}
impl<T, D> SliceInfo<T, D>
where
D: Dimension,
{
#[doc(hidden)]
pub unsafe fn new_unchecked(indices: T, out_dim: PhantomData<D>) -> SliceInfo<T, D> {
SliceInfo { out_dim, indices }
}
}
impl<T, D> SliceInfo<T, D>
where
T: AsRef<[SliceOrIndex]>,
D: Dimension,
{
pub fn new(indices: T) -> Result<SliceInfo<T, D>, ShapeError> {
if let Some(ndim) = D::NDIM {
if ndim != indices.as_ref().iter().filter(|s| s.is_slice()).count() {
return Err(ShapeError::from_kind(ErrorKind::IncompatibleShape));
}
}
Ok(SliceInfo {
out_dim: PhantomData,
indices,
})
}
}
impl<T: ?Sized, D> SliceInfo<T, D>
where
T: AsRef<[SliceOrIndex]>,
D: Dimension,
{
pub fn out_ndim(&self) -> usize {
D::NDIM.unwrap_or_else(|| {
self.indices
.as_ref()
.iter()
.filter(|s| s.is_slice())
.count()
})
}
}
impl<T, D> AsRef<[SliceOrIndex]> for SliceInfo<T, D>
where
T: AsRef<[SliceOrIndex]>,
D: Dimension,
{
fn as_ref(&self) -> &[SliceOrIndex] {
self.indices.as_ref()
}
}
impl<T, D> AsRef<SliceInfo<[SliceOrIndex], D>> for SliceInfo<T, D>
where
T: AsRef<[SliceOrIndex]>,
D: Dimension,
{
fn as_ref(&self) -> &SliceInfo<[SliceOrIndex], D> {
unsafe {
&*(self.indices.as_ref() as *const [SliceOrIndex]
as *const SliceInfo<[SliceOrIndex], D>)
}
}
}
impl<T, D> Copy for SliceInfo<T, D>
where
T: Copy,
D: Dimension,
{
}
impl<T, D> Clone for SliceInfo<T, D>
where
T: Clone,
D: Dimension,
{
fn clone(&self) -> Self {
SliceInfo {
out_dim: PhantomData,
indices: self.indices.clone(),
}
}
}
#[doc(hidden)]
pub trait SliceNextDim<D1, D2> {
fn next_dim(&self, _: PhantomData<D1>) -> PhantomData<D2>;
}
macro_rules! impl_slicenextdim_equal {
($self:ty) => {
impl<D1: Dimension> SliceNextDim<D1, D1> for $self {
fn next_dim(&self, _: PhantomData<D1>) -> PhantomData<D1> {
PhantomData
}
}
};
}
impl_slicenextdim_equal!(isize);
impl_slicenextdim_equal!(usize);
impl_slicenextdim_equal!(i32);
macro_rules! impl_slicenextdim_larger {
(($($generics:tt)*), $self:ty) => {
impl<D1: Dimension, $($generics)*> SliceNextDim<D1, D1::Larger> for $self {
fn next_dim(&self, _: PhantomData<D1>) -> PhantomData<D1::Larger> {
PhantomData
}
}
}
}
impl_slicenextdim_larger!((T), Range<T>);
impl_slicenextdim_larger!((T), RangeInclusive<T>);
impl_slicenextdim_larger!((T), RangeFrom<T>);
impl_slicenextdim_larger!((T), RangeTo<T>);
impl_slicenextdim_larger!((T), RangeToInclusive<T>);
impl_slicenextdim_larger!((), RangeFull);
impl_slicenextdim_larger!((), Slice);
#[macro_export]
macro_rules! s(
(@parse $dim:expr, [$($stack:tt)*] $r:expr;$s:expr) => {
match $r {
r => {
let out_dim = $crate::SliceNextDim::next_dim(&r, $dim);
#[allow(unsafe_code)]
unsafe {
$crate::SliceInfo::new_unchecked(
[$($stack)* $crate::s!(@convert r, $s)],
out_dim,
)
}
}
}
};
(@parse $dim:expr, [$($stack:tt)*] $r:expr) => {
match $r {
r => {
let out_dim = $crate::SliceNextDim::next_dim(&r, $dim);
#[allow(unsafe_code)]
unsafe {
$crate::SliceInfo::new_unchecked(
[$($stack)* $crate::s!(@convert r)],
out_dim,
)
}
}
}
};
(@parse $dim:expr, [$($stack:tt)*] $r:expr;$s:expr ,) => {
$crate::s![@parse $dim, [$($stack)*] $r;$s]
};
(@parse $dim:expr, [$($stack:tt)*] $r:expr ,) => {
$crate::s![@parse $dim, [$($stack)*] $r]
};
(@parse $dim:expr, [$($stack:tt)*] $r:expr;$s:expr, $($t:tt)*) => {
match $r {
r => {
$crate::s![@parse
$crate::SliceNextDim::next_dim(&r, $dim),
[$($stack)* $crate::s!(@convert r, $s),]
$($t)*
]
}
}
};
(@parse $dim:expr, [$($stack:tt)*] $r:expr, $($t:tt)*) => {
match $r {
r => {
$crate::s![@parse
$crate::SliceNextDim::next_dim(&r, $dim),
[$($stack)* $crate::s!(@convert r),]
$($t)*
]
}
}
};
(@parse ::std::marker::PhantomData::<$crate::Ix0>, []) => {
{
#[allow(unsafe_code)]
unsafe {
$crate::SliceInfo::new_unchecked([], ::std::marker::PhantomData::<$crate::Ix0>)
}
}
};
(@parse $($t:tt)*) => { compile_error!("Invalid syntax in s![] call.") };
(@convert $r:expr) => {
<$crate::SliceOrIndex as ::std::convert::From<_>>::from($r)
};
(@convert $r:expr, $s:expr) => {
<$crate::SliceOrIndex as ::std::convert::From<_>>::from($r).step_by($s as isize)
};
($($t:tt)*) => {
&*&$crate::s![@parse ::std::marker::PhantomData::<$crate::Ix0>, [] $($t)*]
};
);
pub trait MultiSlice<'a, A, D>
where
A: 'a,
D: Dimension,
{
type Output;
fn multi_slice_move(&self, view: ArrayViewMut<'a, A, D>) -> Self::Output;
}
impl<'a, A, D> MultiSlice<'a, A, D> for ()
where
A: 'a,
D: Dimension,
{
type Output = ();
fn multi_slice_move(&self, _view: ArrayViewMut<'a, A, D>) -> Self::Output {}
}
impl<'a, A, D, Do0> MultiSlice<'a, A, D> for (&SliceInfo<D::SliceArg, Do0>,)
where
A: 'a,
D: Dimension,
Do0: Dimension,
{
type Output = (ArrayViewMut<'a, A, Do0>,);
fn multi_slice_move(&self, view: ArrayViewMut<'a, A, D>) -> Self::Output {
(view.slice_move(self.0),)
}
}
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,)*> MultiSlice<'a, A, D> for ($(&SliceInfo<D::SliceArg, $all>,)*)
where
A: 'a,
D: Dimension,
$($all: Dimension,)*
{
type Output = ($(ArrayViewMut<'a, A, $all>,)*);
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(),
)
}
}
}
};
(@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!([Do0] Do1);
impl_multislice_tuple!([Do0 Do1] Do2);
impl_multislice_tuple!([Do0 Do1 Do2] Do3);
impl_multislice_tuple!([Do0 Do1 Do2 Do3] Do4);
impl_multislice_tuple!([Do0 Do1 Do2 Do3 Do4] Do5);
impl<'a, A, D, T> MultiSlice<'a, A, D> for &T
where
A: 'a,
D: Dimension,
T: MultiSlice<'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)
}
}