ndarray/dimension/
ndindex.rs

1use std::fmt::Debug;
2
3use super::{stride_offset, stride_offset_checked};
4use crate::itertools::zip;
5use crate::{Dim, Dimension, IntoDimension, Ix, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn, IxDynImpl};
6
7/// Tuple or fixed size arrays that can be used to index an array.
8///
9/// ```
10/// use ndarray::arr2;
11///
12/// let mut a = arr2(&[[0, 1],
13///                    [2, 3]]);
14/// assert_eq!(a[[0, 1]], 1);
15/// assert_eq!(a[[1, 1]], 3);
16/// a[[1, 1]] += 1;
17/// assert_eq!(a[(1, 1)], 4);
18/// ```
19#[allow(clippy::missing_safety_doc)] // TODO: Add doc
20pub unsafe trait NdIndex<E>: Debug
21{
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 D: Dimension
30{
31    fn index_checked(&self, dim: &D, strides: &D) -> Option<isize>
32    {
33        dim.stride_offset_checked(strides, self)
34    }
35    fn index_unchecked(&self, strides: &D) -> isize
36    {
37        D::stride_offset(self, strides)
38    }
39}
40
41unsafe impl NdIndex<Ix0> for ()
42{
43    #[inline]
44    fn index_checked(&self, dim: &Ix0, strides: &Ix0) -> Option<isize>
45    {
46        dim.stride_offset_checked(strides, &Ix0())
47    }
48    #[inline(always)]
49    fn index_unchecked(&self, _strides: &Ix0) -> isize
50    {
51        0
52    }
53}
54
55unsafe impl NdIndex<Ix2> for (Ix, Ix)
56{
57    #[inline]
58    fn index_checked(&self, dim: &Ix2, strides: &Ix2) -> Option<isize>
59    {
60        dim.stride_offset_checked(strides, &Ix2(self.0, self.1))
61    }
62    #[inline]
63    fn index_unchecked(&self, strides: &Ix2) -> isize
64    {
65        stride_offset(self.0, get!(strides, 0)) + stride_offset(self.1, get!(strides, 1))
66    }
67}
68unsafe impl NdIndex<Ix3> for (Ix, Ix, Ix)
69{
70    #[inline]
71    fn index_checked(&self, dim: &Ix3, strides: &Ix3) -> Option<isize>
72    {
73        dim.stride_offset_checked(strides, &self.into_dimension())
74    }
75
76    #[inline]
77    fn index_unchecked(&self, strides: &Ix3) -> isize
78    {
79        stride_offset(self.0, get!(strides, 0))
80            + stride_offset(self.1, get!(strides, 1))
81            + stride_offset(self.2, get!(strides, 2))
82    }
83}
84
85unsafe impl NdIndex<Ix4> for (Ix, Ix, Ix, Ix)
86{
87    #[inline]
88    fn index_checked(&self, dim: &Ix4, strides: &Ix4) -> Option<isize>
89    {
90        dim.stride_offset_checked(strides, &self.into_dimension())
91    }
92    #[inline]
93    fn index_unchecked(&self, strides: &Ix4) -> isize
94    {
95        zip(strides.ix(), self.into_dimension().ix())
96            .map(|(&s, &i)| stride_offset(i, s))
97            .sum()
98    }
99}
100unsafe impl NdIndex<Ix5> for (Ix, Ix, Ix, Ix, Ix)
101{
102    #[inline]
103    fn index_checked(&self, dim: &Ix5, strides: &Ix5) -> Option<isize>
104    {
105        dim.stride_offset_checked(strides, &self.into_dimension())
106    }
107    #[inline]
108    fn index_unchecked(&self, strides: &Ix5) -> isize
109    {
110        zip(strides.ix(), self.into_dimension().ix())
111            .map(|(&s, &i)| stride_offset(i, s))
112            .sum()
113    }
114}
115
116unsafe impl NdIndex<Ix6> for (Ix, Ix, Ix, Ix, Ix, Ix)
117{
118    #[inline]
119    fn index_checked(&self, dim: &Ix6, strides: &Ix6) -> Option<isize>
120    {
121        dim.stride_offset_checked(strides, &self.into_dimension())
122    }
123    #[inline]
124    fn index_unchecked(&self, strides: &Ix6) -> isize
125    {
126        zip(strides.ix(), self.into_dimension().ix())
127            .map(|(&s, &i)| stride_offset(i, s))
128            .sum()
129    }
130}
131
132unsafe impl NdIndex<Ix1> for Ix
133{
134    #[inline]
135    fn index_checked(&self, dim: &Ix1, strides: &Ix1) -> Option<isize>
136    {
137        dim.stride_offset_checked(strides, &Ix1(*self))
138    }
139    #[inline(always)]
140    fn index_unchecked(&self, strides: &Ix1) -> isize
141    {
142        stride_offset(*self, get!(strides, 0))
143    }
144}
145
146unsafe impl NdIndex<IxDyn> for Ix
147{
148    #[inline]
149    fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize>
150    {
151        debug_assert_eq!(dim.ndim(), 1);
152        stride_offset_checked(dim.ix(), strides.ix(), &[*self])
153    }
154    #[inline(always)]
155    fn index_unchecked(&self, strides: &IxDyn) -> isize
156    {
157        debug_assert_eq!(strides.ndim(), 1);
158        stride_offset(*self, get!(strides, 0))
159    }
160}
161
162macro_rules! ndindex_with_array {
163    ($([$n:expr, $ix_n:ident $($index:tt)*])+) => {
164        $(
165        // implement NdIndex<Ix2> for [Ix; 2] and so on
166        unsafe impl NdIndex<$ix_n> for [Ix; $n] {
167            #[inline]
168            fn index_checked(&self, dim: &$ix_n, strides: &$ix_n) -> Option<isize> {
169                dim.stride_offset_checked(strides, &self.into_dimension())
170            }
171
172            #[inline]
173            fn index_unchecked(&self, _strides: &$ix_n) -> isize {
174                $(
175                stride_offset(self[$index], get!(_strides, $index)) +
176                )*
177                0
178            }
179        }
180        )+
181    };
182}
183
184ndindex_with_array! {
185    [0, Ix0]
186    [1, Ix1 0]
187    [2, Ix2 0 1]
188    [3, Ix3 0 1 2]
189    [4, Ix4 0 1 2 3]
190    [5, Ix5 0 1 2 3 4]
191    [6, Ix6 0 1 2 3 4 5]
192}
193
194// implement NdIndex<IxDyn> for Dim<[Ix; 2]> and so on
195unsafe impl<const N: usize> NdIndex<IxDyn> for Dim<[Ix; N]>
196{
197    #[inline]
198    fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize>
199    {
200        debug_assert_eq!(
201            strides.ndim(),
202            N,
203            "Attempted to index with {:?} in array with {} axes",
204            self,
205            strides.ndim()
206        );
207        stride_offset_checked(dim.ix(), strides.ix(), self.ix())
208    }
209
210    #[inline]
211    fn index_unchecked(&self, strides: &IxDyn) -> isize
212    {
213        debug_assert_eq!(
214            strides.ndim(),
215            N,
216            "Attempted to index with {:?} in array with {} axes",
217            self,
218            strides.ndim()
219        );
220        (0..N)
221            .map(|i| stride_offset(get!(self, i), get!(strides, i)))
222            .sum()
223    }
224}
225
226// implement NdIndex<IxDyn> for [Ix; 2] and so on
227unsafe impl<const N: usize> NdIndex<IxDyn> for [Ix; N]
228{
229    #[inline]
230    fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize>
231    {
232        debug_assert_eq!(
233            strides.ndim(),
234            N,
235            "Attempted to index with {:?} in array with {} axes",
236            self,
237            strides.ndim()
238        );
239        stride_offset_checked(dim.ix(), strides.ix(), self)
240    }
241
242    #[inline]
243    fn index_unchecked(&self, strides: &IxDyn) -> isize
244    {
245        debug_assert_eq!(
246            strides.ndim(),
247            N,
248            "Attempted to index with {:?} in array with {} axes",
249            self,
250            strides.ndim()
251        );
252        (0..N)
253            .map(|i| stride_offset(self[i], get!(strides, i)))
254            .sum()
255    }
256}
257
258impl<'a> IntoDimension for &'a [Ix]
259{
260    type Dim = IxDyn;
261    fn into_dimension(self) -> Self::Dim
262    {
263        Dim(IxDynImpl::from(self))
264    }
265}
266
267unsafe impl<'a> NdIndex<IxDyn> for &'a IxDyn
268{
269    fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize>
270    {
271        (**self).index_checked(dim, strides)
272    }
273    fn index_unchecked(&self, strides: &IxDyn) -> isize
274    {
275        (**self).index_unchecked(strides)
276    }
277}
278
279unsafe impl<'a> NdIndex<IxDyn> for &'a [Ix]
280{
281    fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize>
282    {
283        stride_offset_checked(dim.ix(), strides.ix(), self)
284    }
285    fn index_unchecked(&self, strides: &IxDyn) -> isize
286    {
287        zip(strides.ix(), *self)
288            .map(|(&s, &i)| stride_offset(i, s))
289            .sum()
290    }
291}
ndarray/dimension/ndindex.rs

ndarray/dimension/
ndindex.rs
