ndarray/dimension/
ndindex.rs

1use std::fmt::Debug;
2
3use super::{stride_offset, stride_offset_checked};
4use crate::itertools::zip;
5use crate::{
6    Dim, Dimension, IntoDimension, Ix, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn, IxDynImpl,
7};
8
9/// Tuple or fixed size arrays that can be used to index an array.
10///
11/// ```
12/// use ndarray::arr2;
13///
14/// let mut a = arr2(&[[0, 1],
15///                    [2, 3]]);
16/// assert_eq!(a[[0, 1]], 1);
17/// assert_eq!(a[[1, 1]], 3);
18/// a[[1, 1]] += 1;
19/// assert_eq!(a[(1, 1)], 4);
20/// ```
21pub unsafe trait NdIndex<E>: Debug {
22    #[doc(hidden)]
23    fn index_checked(&self, dim: &E, strides: &E) -> Option<isize>;
24    #[doc(hidden)]
25    fn index_unchecked(&self, strides: &E) -> isize;
26}
27
28unsafe impl<D> NdIndex<D> for D
29where
30    D: Dimension,
31{
32    fn index_checked(&self, dim: &D, strides: &D) -> Option<isize> {
33        dim.stride_offset_checked(strides, self)
34    }
35    fn index_unchecked(&self, strides: &D) -> isize {
36        D::stride_offset(self, strides)
37    }
38}
39
40unsafe impl NdIndex<Ix0> for () {
41    #[inline]
42    fn index_checked(&self, dim: &Ix0, strides: &Ix0) -> Option<isize> {
43        dim.stride_offset_checked(strides, &Ix0())
44    }
45    #[inline(always)]
46    fn index_unchecked(&self, _strides: &Ix0) -> isize {
47        0
48    }
49}
50
51unsafe impl NdIndex<Ix2> for (Ix, Ix) {
52    #[inline]
53    fn index_checked(&self, dim: &Ix2, strides: &Ix2) -> Option<isize> {
54        dim.stride_offset_checked(strides, &Ix2(self.0, self.1))
55    }
56    #[inline]
57    fn index_unchecked(&self, strides: &Ix2) -> isize {
58        stride_offset(self.0, get!(strides, 0)) + stride_offset(self.1, get!(strides, 1))
59    }
60}
61unsafe impl NdIndex<Ix3> for (Ix, Ix, Ix) {
62    #[inline]
63    fn index_checked(&self, dim: &Ix3, strides: &Ix3) -> Option<isize> {
64        dim.stride_offset_checked(strides, &self.into_dimension())
65    }
66
67    #[inline]
68    fn index_unchecked(&self, strides: &Ix3) -> isize {
69        stride_offset(self.0, get!(strides, 0))
70            + stride_offset(self.1, get!(strides, 1))
71            + stride_offset(self.2, get!(strides, 2))
72    }
73}
74
75unsafe impl NdIndex<Ix4> for (Ix, Ix, Ix, Ix) {
76    #[inline]
77    fn index_checked(&self, dim: &Ix4, strides: &Ix4) -> Option<isize> {
78        dim.stride_offset_checked(strides, &self.into_dimension())
79    }
80    #[inline]
81    fn index_unchecked(&self, strides: &Ix4) -> isize {
82        zip(strides.ix(), self.into_dimension().ix())
83            .map(|(&s, &i)| stride_offset(i, s))
84            .sum()
85    }
86}
87unsafe impl NdIndex<Ix5> for (Ix, Ix, Ix, Ix, Ix) {
88    #[inline]
89    fn index_checked(&self, dim: &Ix5, strides: &Ix5) -> Option<isize> {
90        dim.stride_offset_checked(strides, &self.into_dimension())
91    }
92    #[inline]
93    fn index_unchecked(&self, strides: &Ix5) -> isize {
94        zip(strides.ix(), self.into_dimension().ix())
95            .map(|(&s, &i)| stride_offset(i, s))
96            .sum()
97    }
98}
99
100unsafe impl NdIndex<Ix1> for Ix {
101    #[inline]
102    fn index_checked(&self, dim: &Ix1, strides: &Ix1) -> Option<isize> {
103        dim.stride_offset_checked(strides, &Ix1(*self))
104    }
105    #[inline(always)]
106    fn index_unchecked(&self, strides: &Ix1) -> isize {
107        stride_offset(*self, get!(strides, 0))
108    }
109}
110
111unsafe impl NdIndex<IxDyn> for Ix {
112    #[inline]
113    fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize> {
114        debug_assert_eq!(dim.ndim(), 1);
115        stride_offset_checked(dim.ix(), strides.ix(), &[*self])
116    }
117    #[inline(always)]
118    fn index_unchecked(&self, strides: &IxDyn) -> isize {
119        debug_assert_eq!(strides.ndim(), 1);
120        stride_offset(*self, get!(strides, 0))
121    }
122}
123
124macro_rules! ndindex_with_array {
125    ($([$n:expr, $ix_n:ident $($index:tt)*])+) => {
126        $(
127        // implement NdIndex<Ix2> for [Ix; 2] and so on
128        unsafe impl NdIndex<$ix_n> for [Ix; $n] {
129            #[inline]
130            fn index_checked(&self, dim: &$ix_n, strides: &$ix_n) -> Option<isize> {
131                dim.stride_offset_checked(strides, &self.into_dimension())
132            }
133
134            #[inline]
135            fn index_unchecked(&self, _strides: &$ix_n) -> isize {
136                $(
137                stride_offset(self[$index], get!(_strides, $index)) +
138                )*
139                0
140            }
141        }
142
143        // implement NdIndex<IxDyn> for Dim<[Ix; 2]> and so on
144        unsafe impl NdIndex<IxDyn> for Dim<[Ix; $n]> {
145            #[inline]
146            fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize> {
147                debug_assert_eq!(strides.ndim(), $n,
148                              "Attempted to index with {:?} in array with {} axes",
149                              self, strides.ndim());
150                stride_offset_checked(dim.ix(), strides.ix(), self.ix())
151            }
152
153            #[inline]
154            fn index_unchecked(&self, strides: &IxDyn) -> isize {
155                debug_assert_eq!(strides.ndim(), $n,
156                              "Attempted to index with {:?} in array with {} axes",
157                              self, strides.ndim());
158                $(
159                stride_offset(get!(self, $index), get!(strides, $index)) +
160                )*
161                0
162            }
163        }
164
165        // implement NdIndex<IxDyn> for [Ix; 2] and so on
166        unsafe impl NdIndex<IxDyn> for [Ix; $n] {
167            #[inline]
168            fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize> {
169                debug_assert_eq!(strides.ndim(), $n,
170                              "Attempted to index with {:?} in array with {} axes",
171                              self, strides.ndim());
172                stride_offset_checked(dim.ix(), strides.ix(), self)
173            }
174
175            #[inline]
176            fn index_unchecked(&self, strides: &IxDyn) -> isize {
177                debug_assert_eq!(strides.ndim(), $n,
178                              "Attempted to index with {:?} in array with {} axes",
179                              self, strides.ndim());
180                $(
181                stride_offset(self[$index], get!(strides, $index)) +
182                )*
183                0
184            }
185        }
186        )+
187    };
188}
189
190ndindex_with_array! {
191    [0, Ix0]
192    [1, Ix1 0]
193    [2, Ix2 0 1]
194    [3, Ix3 0 1 2]
195    [4, Ix4 0 1 2 3]
196    [5, Ix5 0 1 2 3 4]
197    [6, Ix6 0 1 2 3 4 5]
198}
199
200impl<'a> IntoDimension for &'a [Ix] {
201    type Dim = IxDyn;
202    fn into_dimension(self) -> Self::Dim {
203        Dim(IxDynImpl::from(self))
204    }
205}
206
207unsafe impl<'a> NdIndex<IxDyn> for &'a IxDyn {
208    fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize> {
209        (**self).index_checked(dim, strides)
210    }
211    fn index_unchecked(&self, strides: &IxDyn) -> isize {
212        (**self).index_unchecked(strides)
213    }
214}
215
216unsafe impl<'a> NdIndex<IxDyn> for &'a [Ix] {
217    fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize> {
218        stride_offset_checked(dim.ix(), strides.ix(), *self)
219    }
220    fn index_unchecked(&self, strides: &IxDyn) -> isize {
221        zip(strides.ix(), *self)
222            .map(|(&s, &i)| stride_offset(i, s))
223            .sum()
224    }
225}
ndarray/dimension/ndindex.rs

ndarray/dimension/
ndindex.rs
