use std::fmt::Debug;
use std::ops::{Index, IndexMut};
use std::ops::{Add, Sub, Mul, AddAssign, SubAssign, MulAssign};
use itertools::{enumerate, zip};
use {Ix, Ixs, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn, Dim, SliceOrIndex, IxDynImpl};
use IntoDimension;
use RemoveAxis;
use {ArrayView1, ArrayViewMut1};
use Axis;
use super::{
stride_offset,
stride_offset_checked,
};
use super::conversion::Convert;
use super::axes_of;
pub trait Dimension : Clone + Eq + Debug + Send + Sync + Default +
IndexMut<usize, Output=usize> +
Add<Self, Output=Self> +
AddAssign + for<'x> AddAssign<&'x Self> +
Sub<Self, Output=Self> +
SubAssign + for<'x> SubAssign<&'x Self> +
Mul<usize, Output=Self> + Mul<Self, Output=Self> +
MulAssign + for<'x> MulAssign<&'x Self> + MulAssign<usize>
{
const NDIM: Option<usize>;
type SliceArg: ?Sized + AsRef<[SliceOrIndex]>;
type Pattern: IntoDimension<Dim=Self>;
type Smaller: Dimension;
type Larger: Dimension;
fn ndim(&self) -> usize;
fn into_pattern(self) -> Self::Pattern;
fn size(&self) -> usize {
self.slice().iter().fold(1, |s, &a| s * a as usize)
}
fn size_checked(&self) -> Option<usize> {
self.slice().iter().fold(Some(1), |s, &a| s.and_then(|s_| s_.checked_mul(a)))
}
#[doc(hidden)]
fn slice(&self) -> &[Ix];
#[doc(hidden)]
fn slice_mut(&mut self) -> &mut [Ix];
fn as_array_view(&self) -> ArrayView1<Ix> {
ArrayView1::from(self.slice())
}
fn as_array_view_mut(&mut self) -> ArrayViewMut1<Ix> {
ArrayViewMut1::from(self.slice_mut())
}
#[doc(hidden)]
fn equal(&self, rhs: &Self) -> bool {
self.slice() == rhs.slice()
}
#[doc(hidden)]
fn default_strides(&self) -> Self {
let mut strides = Self::zeros(self.ndim());
if self.slice().iter().all(|&d| d != 0) {
let mut it = strides.slice_mut().iter_mut().rev();
while let Some(rs) = it.next() {
*rs = 1;
break;
}
let mut cum_prod = 1;
for (rs, dim) in it.zip(self.slice().iter().rev()) {
cum_prod *= *dim;
*rs = cum_prod;
}
}
strides
}
#[doc(hidden)]
fn fortran_strides(&self) -> Self {
let mut strides = Self::zeros(self.ndim());
if self.slice().iter().all(|&d| d != 0) {
let mut it = strides.slice_mut().iter_mut();
while let Some(rs) = it.next() {
*rs = 1;
break;
}
let mut cum_prod = 1;
for (rs, dim) in it.zip(self.slice()) {
cum_prod *= *dim;
*rs = cum_prod;
}
}
strides
}
fn zeros(ndim: usize) -> Self;
#[doc(hidden)]
#[inline]
fn first_index(&self) -> Option<Self> {
for ax in self.slice().iter() {
if *ax == 0 {
return None;
}
}
Some(Self::zeros(self.ndim()))
}
#[doc(hidden)]
#[inline]
fn next_for(&self, index: Self) -> Option<Self> {
let mut index = index;
let mut done = false;
for (&dim, ix) in zip(self.slice(), index.slice_mut()).rev() {
*ix += 1;
if *ix == dim {
*ix = 0;
} else {
done = true;
break;
}
}
if done {
Some(index)
} else {
None
}
}
#[doc(hidden)]
#[inline]
fn next_for_f(&self, index: &mut Self) -> bool {
let mut end_iteration = true;
for (&dim, ix) in zip(self.slice(), index.slice_mut()) {
*ix += 1;
if *ix == dim {
*ix = 0;
} else {
end_iteration = false;
break;
}
}
!end_iteration
}
#[doc(hidden)]
fn stride_offset(index: &Self, strides: &Self) -> isize {
let mut offset = 0;
for (&i, &s) in izip!(index.slice(), strides.slice()) {
offset += stride_offset(i, s);
}
offset
}
#[doc(hidden)]
fn stride_offset_checked(&self, strides: &Self, index: &Self) -> Option<isize> {
stride_offset_checked(self.slice(), strides.slice(), index.slice())
}
#[doc(hidden)]
fn last_elem(&self) -> usize {
if self.ndim() == 0 { 0 } else { self.slice()[self.ndim() - 1] }
}
#[doc(hidden)]
fn set_last_elem(&mut self, i: usize) {
let nd = self.ndim();
self.slice_mut()[nd - 1] = i;
}
#[doc(hidden)]
fn is_contiguous(dim: &Self, strides: &Self) -> bool {
let defaults = dim.default_strides();
if strides.equal(&defaults) {
return true;
}
if dim.ndim() == 1 { return false; }
let order = strides._fastest_varying_stride_order();
let strides = strides.slice();
let dim_slice = dim.slice();
let mut cstride = 1;
for &i in order.slice() {
if dim_slice[i] != 1 && strides[i] != cstride {
return false;
}
cstride *= dim_slice[i];
}
true
}
#[doc(hidden)]
fn _fastest_varying_stride_order(&self) -> Self {
let mut indices = self.clone();
for (i, elt) in enumerate(indices.slice_mut()) {
*elt = i;
}
let strides = self.slice();
indices.slice_mut().sort_by_key(|&i| strides[i]);
indices
}
#[doc(hidden)]
fn min_stride_axis(&self, strides: &Self) -> Axis {
let n = match self.ndim() {
0 => panic!("min_stride_axis: Array must have ndim > 0"),
1 => return Axis(0),
n => n,
};
axes_of(self, strides)
.rev()
.min_by_key(|ax| ax.stride().abs())
.map_or(Axis(n - 1), |ax| ax.axis())
}
#[doc(hidden)]
fn max_stride_axis(&self, strides: &Self) -> Axis {
match self.ndim() {
0 => panic!("max_stride_axis: Array must have ndim > 0"),
1 => return Axis(0),
_ => { }
}
axes_of(self, strides)
.filter(|ax| ax.len() > 1)
.max_by_key(|ax| ax.stride().abs())
.map_or(Axis(0), |ax| ax.axis())
}
fn into_dyn(self) -> IxDyn {
IxDyn(self.slice())
}
#[doc(hidden)]
fn from_dimension<D2: Dimension>(d: &D2) -> Option<Self> {
let mut s = Self::default();
if s.ndim() == d.ndim() {
for i in 0..d.ndim() {
s[i] = d[i];
}
Some(s)
} else {
None
}
}
#[doc(hidden)]
fn insert_axis(&self, axis: Axis) -> Self::Larger;
#[doc(hidden)]
fn try_remove_axis(&self, axis: Axis) -> Self::Smaller;
private_decl!{}
}
macro_rules! impl_insert_axis_array(
($n:expr) => (
fn insert_axis(&self, axis: Axis) -> Self::Larger {
debug_assert!(axis.index() <= $n);
let mut out = [1; $n + 1];
out[0..axis.index()].copy_from_slice(&self.slice()[0..axis.index()]);
out[axis.index()+1..$n+1].copy_from_slice(&self.slice()[axis.index()..$n]);
Dim(out)
}
);
);
impl Dimension for Dim<[Ix; 0]> {
const NDIM: Option<usize> = Some(0);
type SliceArg = [SliceOrIndex; 0];
type Pattern = ();
type Smaller = Self;
type Larger = Ix1;
#[inline]
fn ndim(&self) -> usize { 0 }
#[inline]
fn slice(&self) -> &[Ix] { &[] }
#[inline]
fn slice_mut(&mut self) -> &mut [Ix] { &mut [] }
#[inline]
fn _fastest_varying_stride_order(&self) -> Self { Ix0() }
#[inline]
fn into_pattern(self) -> Self::Pattern { }
#[inline]
fn zeros(ndim: usize) -> Self {
assert_eq!(ndim, 0);
Self::default()
}
#[inline]
fn next_for(&self, _index: Self) -> Option<Self> {
None
}
#[inline]
impl_insert_axis_array!(0);
#[inline]
fn try_remove_axis(&self, _ignore: Axis) -> Self::Smaller {
*self
}
private_impl!{}
}
impl Dimension for Dim<[Ix; 1]> {
const NDIM: Option<usize> = Some(1);
type SliceArg = [SliceOrIndex; 1];
type Pattern = Ix;
type Smaller = Ix0;
type Larger = Ix2;
#[inline]
fn ndim(&self) -> usize { 1 }
#[inline]
fn slice(&self) -> &[Ix] { self.ix() }
#[inline]
fn slice_mut(&mut self) -> &mut [Ix] { self.ixm() }
#[inline]
fn into_pattern(self) -> Self::Pattern {
get!(&self, 0)
}
#[inline]
fn zeros(ndim: usize) -> Self {
assert_eq!(ndim, 1);
Self::default()
}
#[inline]
fn next_for(&self, mut index: Self) -> Option<Self> {
getm!(index, 0) += 1;
if get!(&index, 0) < get!(self, 0) {
Some(index)
} else {
None
}
}
#[inline]
fn equal(&self, rhs: &Self) -> bool {
get!(self, 0) == get!(rhs, 0)
}
#[inline]
fn size(&self) -> usize { get!(self, 0) }
#[inline]
fn size_checked(&self) -> Option<usize> { Some(get!(self, 0)) }
#[inline]
fn default_strides(&self) -> Self {
if get!(self, 0) == 0 {
Ix1(0)
} else {
Ix1(1)
}
}
#[inline]
fn _fastest_varying_stride_order(&self) -> Self {
Ix1(0)
}
#[inline(always)]
fn min_stride_axis(&self, _: &Self) -> Axis {
Axis(0)
}
#[inline(always)]
fn max_stride_axis(&self, _: &Self) -> Axis {
Axis(0)
}
#[inline]
fn first_index(&self) -> Option<Self> {
if get!(self, 0) != 0 {
Some(Ix1(0))
} else {
None
}
}
#[inline(always)]
fn stride_offset(index: &Self, stride: &Self) -> isize {
stride_offset(get!(index, 0), get!(stride, 0))
}
#[inline]
fn stride_offset_checked(&self, stride: &Self, index: &Self) -> Option<isize> {
if get!(index, 0) < get!(self, 0) {
Some(stride_offset(get!(index, 0), get!(stride, 0)))
} else {
None
}
}
#[inline]
impl_insert_axis_array!(1);
#[inline]
fn try_remove_axis(&self, axis: Axis) -> Self::Smaller {
self.remove_axis(axis)
}
private_impl!{}
}
impl Dimension for Dim<[Ix; 2]> {
const NDIM: Option<usize> = Some(2);
type SliceArg = [SliceOrIndex; 2];
type Pattern = (Ix, Ix);
type Smaller = Ix1;
type Larger = Ix3;
#[inline]
fn ndim(&self) -> usize { 2 }
#[inline]
fn into_pattern(self) -> Self::Pattern {
self.ix().convert()
}
#[inline]
fn slice(&self) -> &[Ix] { self.ix() }
#[inline]
fn slice_mut(&mut self) -> &mut [Ix] { self.ixm() }
#[inline]
fn zeros(ndim: usize) -> Self {
assert_eq!(ndim, 2);
Self::default()
}
#[inline]
fn next_for(&self, index: Self) -> Option<Self> {
let mut i = get!(&index, 0);
let mut j = get!(&index, 1);
let imax = get!(self, 0);
let jmax = get!(self, 1);
j += 1;
if j >= jmax {
j = 0;
i += 1;
if i >= imax {
return None;
}
}
Some(Ix2(i, j))
}
#[inline]
fn equal(&self, rhs: &Self) -> bool {
get!(self, 0) == get!(rhs, 0) && get!(self, 1) == get!(rhs, 1)
}
#[inline]
fn size(&self) -> usize { get!(self, 0) * get!(self, 1) }
#[inline]
fn size_checked(&self) -> Option<usize> {
let m = get!(self, 0);
let n = get!(self, 1);
(m as usize).checked_mul(n as usize)
}
#[inline]
fn last_elem(&self) -> usize {
get!(self, 1)
}
#[inline]
fn set_last_elem(&mut self, i: usize) {
getm!(self, 1) = i;
}
#[inline]
fn default_strides(&self) -> Self {
let m = get!(self, 0);
let n = get!(self, 1);
if m == 0 || n == 0 {
Ix2(0, 0)
} else {
Ix2(n, 1)
}
}
#[inline]
fn fortran_strides(&self) -> Self {
let m = get!(self, 0);
let n = get!(self, 1);
if m == 0 || n == 0 {
Ix2(0, 0)
} else {
Ix2(1, m)
}
}
#[inline]
fn _fastest_varying_stride_order(&self) -> Self {
if get!(self, 0) as Ixs <= get!(self, 1) as Ixs { Ix2(0, 1) } else { Ix2(1, 0) }
}
#[inline]
fn min_stride_axis(&self, strides: &Self) -> Axis {
let s = get!(strides, 0) as Ixs;
let t = get!(strides, 1) as Ixs;
if s.abs() < t.abs() {
Axis(0)
} else {
Axis(1)
}
}
#[inline]
fn first_index(&self) -> Option<Self> {
let m = get!(self, 0);
let n = get!(self, 1);
if m != 0 && n != 0 {
Some(Ix2(0, 0))
} else {
None
}
}
#[inline(always)]
fn stride_offset(index: &Self, strides: &Self) -> isize {
let i = get!(index, 0);
let j = get!(index, 1);
let s = get!(strides, 0);
let t = get!(strides, 1);
stride_offset(i, s) + stride_offset(j, t)
}
#[inline]
fn stride_offset_checked(&self, strides: &Self, index: &Self) -> Option<isize>
{
let m = get!(self, 0);
let n = get!(self, 1);
let i = get!(index, 0);
let j = get!(index, 1);
let s = get!(strides, 0);
let t = get!(strides, 1);
if i < m && j < n {
Some(stride_offset(i, s) + stride_offset(j, t))
} else {
None
}
}
#[inline]
impl_insert_axis_array!(2);
#[inline]
fn try_remove_axis(&self, axis: Axis) -> Self::Smaller {
self.remove_axis(axis)
}
private_impl!{}
}
impl Dimension for Dim<[Ix; 3]> {
const NDIM: Option<usize> = Some(3);
type SliceArg = [SliceOrIndex; 3];
type Pattern = (Ix, Ix, Ix);
type Smaller = Ix2;
type Larger = Ix4;
#[inline]
fn ndim(&self) -> usize { 3 }
#[inline]
fn into_pattern(self) -> Self::Pattern {
self.ix().convert()
}
#[inline]
fn slice(&self) -> &[Ix] { self.ix() }
#[inline]
fn slice_mut(&mut self) -> &mut [Ix] { self.ixm() }
#[inline]
fn size(&self) -> usize {
let m = get!(self, 0);
let n = get!(self, 1);
let o = get!(self, 2);
m as usize * n as usize * o as usize
}
#[inline]
fn zeros(ndim: usize) -> Self {
assert_eq!(ndim, 3);
Self::default()
}
#[inline]
fn next_for(&self, index: Self) -> Option<Self> {
let mut i = get!(&index, 0);
let mut j = get!(&index, 1);
let mut k = get!(&index, 2);
let imax = get!(self, 0);
let jmax = get!(self, 1);
let kmax = get!(self, 2);
k += 1;
if k == kmax {
k = 0;
j += 1;
if j == jmax {
j = 0;
i += 1;
if i == imax {
return None;
}
}
}
Some(Ix3(i, j, k))
}
#[inline]
fn stride_offset(index: &Self, strides: &Self) -> isize {
let i = get!(index, 0);
let j = get!(index, 1);
let k = get!(index, 2);
let s = get!(strides, 0);
let t = get!(strides, 1);
let u = get!(strides, 2);
stride_offset(i, s) + stride_offset(j, t) + stride_offset(k, u)
}
#[inline]
fn stride_offset_checked(&self, strides: &Self, index: &Self) -> Option<isize>
{
let m = get!(self, 0);
let n = get!(self, 1);
let l = get!(self, 2);
let i = get!(index, 0);
let j = get!(index, 1);
let k = get!(index, 2);
let s = get!(strides, 0);
let t = get!(strides, 1);
let u = get!(strides, 2);
if i < m && j < n && k < l {
Some(stride_offset(i, s) + stride_offset(j, t) + stride_offset(k, u))
} else {
None
}
}
#[inline]
fn _fastest_varying_stride_order(&self) -> Self {
let mut stride = *self;
let mut order = Ix3(0, 1, 2);
macro_rules! swap {
($stride:expr, $order:expr, $x:expr, $y:expr) => {
if $stride[$x] > $stride[$y] {
$stride.swap($x, $y);
$order.ixm().swap($x, $y);
}
}
}
{
let strides = stride.slice_mut();
swap![strides, order, 1, 2];
swap![strides, order, 0, 1];
swap![strides, order, 1, 2];
}
order
}
#[inline]
impl_insert_axis_array!(3);
#[inline]
fn try_remove_axis(&self, axis: Axis) -> Self::Smaller {
self.remove_axis(axis)
}
private_impl!{}
}
macro_rules! large_dim {
($n:expr, $name:ident, $pattern:ty, $larger:ty, { $($insert_axis:tt)* }) => (
impl Dimension for Dim<[Ix; $n]> {
const NDIM: Option<usize> = Some($n);
type SliceArg = [SliceOrIndex; $n];
type Pattern = $pattern;
type Smaller = Dim<[Ix; $n - 1]>;
type Larger = $larger;
#[inline]
fn ndim(&self) -> usize { $n }
#[inline]
fn into_pattern(self) -> Self::Pattern {
self.ix().convert()
}
#[inline]
fn slice(&self) -> &[Ix] { self.ix() }
#[inline]
fn slice_mut(&mut self) -> &mut [Ix] { self.ixm() }
#[inline]
fn zeros(ndim: usize) -> Self {
assert_eq!(ndim, $n);
Self::default()
}
#[inline]
$($insert_axis)*
#[inline]
fn try_remove_axis(&self, axis: Axis) -> Self::Smaller {
self.remove_axis(axis)
}
private_impl!{}
}
)
}
large_dim!(4, Ix4, (Ix, Ix, Ix, Ix), Ix5, {
impl_insert_axis_array!(4);
});
large_dim!(5, Ix5, (Ix, Ix, Ix, Ix, Ix), Ix6, {
impl_insert_axis_array!(5);
});
large_dim!(6, Ix6, (Ix, Ix, Ix, Ix, Ix, Ix), IxDyn, {
fn insert_axis(&self, axis: Axis) -> Self::Larger {
debug_assert!(axis.index() <= self.ndim());
let mut out = Vec::with_capacity(self.ndim() + 1);
out.extend_from_slice(&self.slice()[0..axis.index()]);
out.push(1);
out.extend_from_slice(&self.slice()[axis.index()..self.ndim()]);
Dim(out)
}
});
impl Dimension for IxDyn
{
const NDIM: Option<usize> = None;
type SliceArg = [SliceOrIndex];
type Pattern = Self;
type Smaller = Self;
type Larger = Self;
#[inline]
fn ndim(&self) -> usize { self.ix().len() }
#[inline]
fn slice(&self) -> &[Ix] { self.ix() }
#[inline]
fn slice_mut(&mut self) -> &mut [Ix] { self.ixm() }
#[inline]
fn into_pattern(self) -> Self::Pattern {
self
}
#[inline]
fn zeros(ndim: usize) -> Self {
IxDyn::zeros(ndim)
}
#[inline]
fn insert_axis(&self, axis: Axis) -> Self::Larger {
debug_assert!(axis.index() <= self.ndim());
Dim::new(self.ix().insert(axis.index()))
}
#[inline]
fn try_remove_axis(&self, axis: Axis) -> Self::Smaller {
if self.ndim() > 0 {
self.remove_axis(axis)
} else {
self.clone()
}
}
fn from_dimension<D2: Dimension>(d: &D2) -> Option<Self> {
Some(IxDyn(d.slice()))
}
private_impl!{}
}
impl Index<usize> for Dim<IxDynImpl>
{
type Output = <IxDynImpl as Index<usize>>::Output;
fn index(&self, index: usize) -> &Self::Output {
&self.ix()[index]
}
}
impl IndexMut<usize> for Dim<IxDynImpl>
{
fn index_mut(&mut self, index: usize) -> &mut Self::Output {
&mut self.ixm()[index]
}
}