/rust/registry/src/index.crates.io-1949cf8c6b5b557f/ndarray-0.16.1/src/dimension/ndindex.rs

Source
use std::fmt::Debug;

use super::{stride_offset, stride_offset_checked};
use crate::itertools::zip;
use crate::{Dim, Dimension, IntoDimension, Ix, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn, IxDynImpl};

/// Tuple or fixed size arrays that can be used to index an array.
///
/// ```
/// use ndarray::arr2;
///
/// let mut a = arr2(&[[0, 1],
///                    [2, 3]]);
/// assert_eq!(a[[0, 1]], 1);
/// assert_eq!(a[[1, 1]], 3);
/// a[[1, 1]] += 1;
/// assert_eq!(a[(1, 1)], 4);
/// ```
#[allow(clippy::missing_safety_doc)] // TODO: Add doc
pub unsafe trait NdIndex<E>: Debug
{
    #[doc(hidden)]
    fn index_checked(&self, dim: &E, strides: &E) -> Option<isize>;
    #[doc(hidden)]
    fn index_unchecked(&self, strides: &E) -> isize;
}

unsafe impl<D> NdIndex<D> for D
where D: Dimension
{
    fn index_checked(&self, dim: &D, strides: &D) -> Option<isize>
    {
        dim.stride_offset_checked(strides, self)
    }
    fn index_unchecked(&self, strides: &D) -> isize
    {
        D::stride_offset(self, strides)
    }
}

unsafe impl NdIndex<Ix0> for ()
{
    #[inline]
    fn index_checked(&self, dim: &Ix0, strides: &Ix0) -> Option<isize>
    {
        dim.stride_offset_checked(strides, &Ix0())
    }
    #[inline(always)]
    fn index_unchecked(&self, _strides: &Ix0) -> isize
    {
        0
    }
}

unsafe impl NdIndex<Ix2> for (Ix, Ix)
{
    #[inline]
    fn index_checked(&self, dim: &Ix2, strides: &Ix2) -> Option<isize>
    {
        dim.stride_offset_checked(strides, &Ix2(self.0, self.1))
    }
    #[inline]
    fn index_unchecked(&self, strides: &Ix2) -> isize
    {
        stride_offset(self.0, get!(strides, 0)) + stride_offset(self.1, get!(strides, 1))
    }
}
unsafe impl NdIndex<Ix3> for (Ix, Ix, Ix)
{
    #[inline]
    fn index_checked(&self, dim: &Ix3, strides: &Ix3) -> Option<isize>
    {
        dim.stride_offset_checked(strides, &self.into_dimension())
    }

    #[inline]
    fn index_unchecked(&self, strides: &Ix3) -> isize
    {
        stride_offset(self.0, get!(strides, 0))
            + stride_offset(self.1, get!(strides, 1))
            + stride_offset(self.2, get!(strides, 2))
    }
}

unsafe impl NdIndex<Ix4> for (Ix, Ix, Ix, Ix)
{
    #[inline]
    fn index_checked(&self, dim: &Ix4, strides: &Ix4) -> Option<isize>
    {
        dim.stride_offset_checked(strides, &self.into_dimension())
    }
    #[inline]
    fn index_unchecked(&self, strides: &Ix4) -> isize
    {
        zip(strides.ix(), self.into_dimension().ix())
            .map(|(&s, &i)| stride_offset(i, s))
            .sum()
    }
}
unsafe impl NdIndex<Ix5> for (Ix, Ix, Ix, Ix, Ix)
{
    #[inline]
    fn index_checked(&self, dim: &Ix5, strides: &Ix5) -> Option<isize>
    {
        dim.stride_offset_checked(strides, &self.into_dimension())
    }
    #[inline]
    fn index_unchecked(&self, strides: &Ix5) -> isize
    {
        zip(strides.ix(), self.into_dimension().ix())
            .map(|(&s, &i)| stride_offset(i, s))
            .sum()
    }
}

unsafe impl NdIndex<Ix6> for (Ix, Ix, Ix, Ix, Ix, Ix)
{
    #[inline]
    fn index_checked(&self, dim: &Ix6, strides: &Ix6) -> Option<isize>
    {
        dim.stride_offset_checked(strides, &self.into_dimension())
    }
    #[inline]
    fn index_unchecked(&self, strides: &Ix6) -> isize
    {
        zip(strides.ix(), self.into_dimension().ix())
            .map(|(&s, &i)| stride_offset(i, s))
            .sum()
    }
}

unsafe impl NdIndex<Ix1> for Ix
{
    #[inline]
    fn index_checked(&self, dim: &Ix1, strides: &Ix1) -> Option<isize>
    {
        dim.stride_offset_checked(strides, &Ix1(*self))
    }
    #[inline(always)]
    fn index_unchecked(&self, strides: &Ix1) -> isize
    {
        stride_offset(*self, get!(strides, 0))
    }
}

unsafe impl NdIndex<IxDyn> for Ix
{
    #[inline]
    fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize>
    {
        debug_assert_eq!(dim.ndim(), 1);
        stride_offset_checked(dim.ix(), strides.ix(), &[*self])
    }
    #[inline(always)]
    fn index_unchecked(&self, strides: &IxDyn) -> isize
    {
        debug_assert_eq!(strides.ndim(), 1);
        stride_offset(*self, get!(strides, 0))
    }
}

macro_rules! ndindex_with_array {
    ($([$n:expr, $ix_n:ident $($index:tt)*])+) => {
        $(
        // implement NdIndex<Ix2> for [Ix; 2] and so on
        unsafe impl NdIndex<$ix_n> for [Ix; $n] {
            #[inline]
            fn index_checked(&self, dim: &$ix_n, strides: &$ix_n) -> Option<isize> {
                dim.stride_offset_checked(strides, &self.into_dimension())
            }

            #[inline]
            fn index_unchecked(&self, _strides: &$ix_n) -> isize {
                $(
                stride_offset(self[$index], get!(_strides, $index)) +
                )*
                0
            }
        }
        )+
    };
}

ndindex_with_array! {
    [0, Ix0]
    [1, Ix1 0]
    [2, Ix2 0 1]
    [3, Ix3 0 1 2]
    [4, Ix4 0 1 2 3]
    [5, Ix5 0 1 2 3 4]
    [6, Ix6 0 1 2 3 4 5]
}

// implement NdIndex<IxDyn> for Dim<[Ix; 2]> and so on
unsafe impl<const N: usize> NdIndex<IxDyn> for Dim<[Ix; N]>
{
    #[inline]
    fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize>
    {
        debug_assert_eq!(
            strides.ndim(),
            N,
            "Attempted to index with {:?} in array with {} axes",
            self,
            strides.ndim()
        );
        stride_offset_checked(dim.ix(), strides.ix(), self.ix())
    }

    #[inline]
    fn index_unchecked(&self, strides: &IxDyn) -> isize
    {
        debug_assert_eq!(
            strides.ndim(),
            N,
            "Attempted to index with {:?} in array with {} axes",
            self,
            strides.ndim()
        );
        (0..N)
            .map(|i| stride_offset(get!(self, i), get!(strides, i)))
            .sum()
    }
}

// implement NdIndex<IxDyn> for [Ix; 2] and so on
unsafe impl<const N: usize> NdIndex<IxDyn> for [Ix; N]
{
    #[inline]
    fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize>
    {
        debug_assert_eq!(
            strides.ndim(),
            N,
            "Attempted to index with {:?} in array with {} axes",
            self,
            strides.ndim()
        );
        stride_offset_checked(dim.ix(), strides.ix(), self)
    }

    #[inline]
    fn index_unchecked(&self, strides: &IxDyn) -> isize
    {
        debug_assert_eq!(
            strides.ndim(),
            N,
            "Attempted to index with {:?} in array with {} axes",
            self,
            strides.ndim()
        );
        (0..N)
            .map(|i| stride_offset(self[i], get!(strides, i)))
            .sum()
    }
}

impl<'a> IntoDimension for &'a [Ix]
{
    type Dim = IxDyn;
    fn into_dimension(self) -> Self::Dim
    {
        Dim(IxDynImpl::from(self))
    }
}

unsafe impl<'a> NdIndex<IxDyn> for &'a IxDyn
{
    fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize>
    {
        (**self).index_checked(dim, strides)
    }
    fn index_unchecked(&self, strides: &IxDyn) -> isize
    {
        (**self).index_unchecked(strides)
    }
}

unsafe impl<'a> NdIndex<IxDyn> for &'a [Ix]
{
    fn index_checked(&self, dim: &IxDyn, strides: &IxDyn) -> Option<isize>
    {
        stride_offset_checked(dim.ix(), strides.ix(), self)
    }
    fn index_unchecked(&self, strides: &IxDyn) -> isize
    {
        zip(strides.ix(), *self)
            .map(|(&s, &i)| stride_offset(i, s))
            .sum()
    }
}

Coverage Report

Created: 2025-11-28 06:44

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