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

ndarray/dimension/
ndindex.rs
