use std::fmt::Debug;
use std::ops::{Add, AddAssign, Mul, MulAssign, Sub, SubAssign};
use std::ops::{Index, IndexMut};
use alloc::vec::Vec;
use super::axes_of;
use super::conversion::Convert;
use super::ops::DimAdd;
use super::{stride_offset, stride_offset_checked};
use crate::itertools::{enumerate, zip};
use crate::{Axis, DimMax};
use crate::IntoDimension;
use crate::RemoveAxis;
use crate::{ArrayView1, ArrayViewMut1};
use crate::{Dim, Ix, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn, IxDynImpl, Ixs};
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>
+ DimMax<Ix0, Output=Self>
+ DimMax<Self, Output=Self>
+ DimMax<IxDyn, Output=IxDyn>
+ DimMax<<Self as Dimension>::Smaller, Output=Self>
+ DimMax<<Self as Dimension>::Larger, Output=<Self as Dimension>::Larger>
+ DimAdd<Self>
+ DimAdd<<Self as Dimension>::Smaller>
+ DimAdd<<Self as Dimension>::Larger>
+ DimAdd<Ix0, Output = Self>
+ DimAdd<Ix1, Output = <Self as Dimension>::Larger>
+ DimAdd<IxDyn, Output = IxDyn>
{
const NDIM: Option<usize>;
type Pattern: IntoDimension<Dim = Self> + Clone + Debug + PartialEq + Eq + Default;
type Smaller: Dimension;
type Larger: Dimension + RemoveAxis;
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();
if let Some(rs) = it.next() {
*rs = 1;
}
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();
if let Some(rs) = it.next() {
*rs = 1;
}
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 strides_equivalent<D>(&self, strides1: &Self, strides2: &D) -> bool
where
D: Dimension,
{
let shape_ndim = self.ndim();
shape_ndim == strides1.ndim()
&& shape_ndim == strides2.ndim()
&& izip!(self.slice(), strides1.slice(), strides2.slice())
.all(|(&d, &s1, &s2)| d <= 1 || s1 as isize == s2 as isize)
}
#[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 strides[0] as isize == -1;
}
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] as isize).abs() as usize != 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] as isize).abs());
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) => (
#[inline]
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].copy_from_slice(&self.slice()[axis.index()..$n]);
Dim(out)
}
);
);
impl Dimension for Dim<[Ix; 0]> {
const NDIM: Option<usize> = Some(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
}
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 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
}
}
impl_insert_axis_array!(1);
#[inline]
fn try_remove_axis(&self, axis: Axis) -> Self::Smaller {
self.remove_axis(axis)
}
fn from_dimension<D2: Dimension>(d: &D2) -> Option<Self> {
if 1 == d.ndim() {
Some(Ix1(d[0]))
} else {
None
}
}
private_impl! {}
}
impl Dimension for Dim<[Ix; 2]> {
const NDIM: Option<usize> = Some(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).abs() <= (get!(self, 1) as Ixs).abs() {
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
}
}
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 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] as isize).abs() > ($stride[$y] as isize).abs() {
$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
}
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 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()
}
$($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 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()))
}
fn into_dyn(self) -> IxDyn {
self
}
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]
}
}