mdarray/index/
slice.rs

1use core::ops::{
2    Bound, Index, IndexMut, Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive,
3};
4
5use crate::index;
6use crate::layout::{Dense, Layout};
7use crate::mapping::Mapping;
8use crate::shape::Shape;
9use crate::slice::Slice;
10
11/// Array slice index trait, for an element or a subslice.
12pub trait SliceIndex<T, S: Shape, L: Layout> {
13    /// Array element or subslice type.
14    type Output: ?Sized;
15
16    #[doc(hidden)]
17    unsafe fn get_unchecked(self, slice: &Slice<T, S, L>) -> &Self::Output;
18
19    #[doc(hidden)]
20    unsafe fn get_unchecked_mut(self, slice: &mut Slice<T, S, L>) -> &mut Self::Output;
21
22    #[doc(hidden)]
23    fn index(self, slice: &Slice<T, S, L>) -> &Self::Output;
24
25    #[doc(hidden)]
26    fn index_mut(self, slice: &mut Slice<T, S, L>) -> &mut Self::Output;
27}
28
29impl<T, S: Shape, L: Layout> SliceIndex<T, S, L> for &[usize] {
30    type Output = T;
31
32    #[inline]
33    unsafe fn get_unchecked(self, slice: &Slice<T, S, L>) -> &T {
34        unsafe { &*slice.as_ptr().offset(slice.mapping().offset(self)) }
35    }
36
37    #[inline]
38    unsafe fn get_unchecked_mut(self, slice: &mut Slice<T, S, L>) -> &mut T {
39        unsafe { &mut *slice.as_mut_ptr().offset(slice.mapping().offset(self)) }
40    }
41
42    #[inline]
43    fn index(self, slice: &Slice<T, S, L>) -> &T {
44        assert!(self.len() == slice.rank(), "invalid rank");
45
46        for i in 0..self.len() {
47            if self[i] >= slice.dim(i) {
48                index::panic_bounds_check(self[i], slice.dim(i));
49            }
50        }
51
52        unsafe { SliceIndex::get_unchecked(self, slice) }
53    }
54
55    #[inline]
56    fn index_mut(self, slice: &mut Slice<T, S, L>) -> &mut T {
57        assert!(self.len() == slice.rank(), "invalid rank");
58
59        for i in 0..self.len() {
60            if self[i] >= slice.dim(i) {
61                index::panic_bounds_check(self[i], slice.dim(i));
62            }
63        }
64
65        unsafe { SliceIndex::get_unchecked_mut(self, slice) }
66    }
67}
68
69impl<T, const N: usize, S: Shape, L: Layout> SliceIndex<T, S, L> for [usize; N] {
70    type Output = T;
71
72    #[inline]
73    unsafe fn get_unchecked(self, slice: &Slice<T, S, L>) -> &T {
74        unsafe { SliceIndex::get_unchecked(&self[..], slice) }
75    }
76
77    #[inline]
78    unsafe fn get_unchecked_mut(self, slice: &mut Slice<T, S, L>) -> &mut T {
79        unsafe { SliceIndex::get_unchecked_mut(&self[..], slice) }
80    }
81
82    #[inline]
83    fn index(self, slice: &Slice<T, S, L>) -> &T {
84        SliceIndex::index(&self[..], slice)
85    }
86
87    #[inline]
88    fn index_mut(self, slice: &mut Slice<T, S, L>) -> &mut T {
89        SliceIndex::index_mut(&self[..], slice)
90    }
91}
92
93impl<T, S: Shape, L: Layout> SliceIndex<T, S, L> for usize {
94    type Output = T;
95
96    #[inline]
97    unsafe fn get_unchecked(self, slice: &Slice<T, S, L>) -> &T {
98        unsafe { &*slice.as_ptr().offset(slice.mapping().linear_offset(self)) }
99    }
100
101    #[inline]
102    unsafe fn get_unchecked_mut(self, slice: &mut Slice<T, S, L>) -> &mut T {
103        unsafe { &mut *slice.as_mut_ptr().offset(slice.mapping().linear_offset(self)) }
104    }
105
106    #[inline]
107    fn index(self, slice: &Slice<T, S, L>) -> &T {
108        if self >= slice.len() {
109            index::panic_bounds_check(self, slice.len());
110        }
111
112        unsafe { SliceIndex::get_unchecked(self, slice) }
113    }
114
115    #[inline]
116    fn index_mut(self, slice: &mut Slice<T, S, L>) -> &mut T {
117        if self >= slice.len() {
118            index::panic_bounds_check(self, slice.len());
119        }
120
121        unsafe { SliceIndex::get_unchecked_mut(self, slice) }
122    }
123}
124
125macro_rules! impl_slice_index {
126    ($type:ty) => {
127        impl<T, S: Shape> SliceIndex<T, S, Dense> for $type {
128            type Output = [T];
129
130            #[inline]
131            unsafe fn get_unchecked(self, slice: &Slice<T, S>) -> &[T] {
132                unsafe { <&[T]>::from(slice.flatten()).get_unchecked(self) }
133            }
134
135            #[inline]
136            unsafe fn get_unchecked_mut(self, slice: &mut Slice<T, S>) -> &mut [T] {
137                unsafe { <&mut [T]>::from(slice.flatten_mut()).get_unchecked_mut(self) }
138            }
139
140            #[inline]
141            fn index(self, slice: &Slice<T, S>) -> &[T] {
142                <&[T]>::from(slice.flatten()).index(self)
143            }
144
145            #[inline]
146            fn index_mut(self, slice: &mut Slice<T, S>) -> &mut [T] {
147                <&mut [T]>::from(slice.flatten_mut()).index_mut(self)
148            }
149        }
150    };
151}
152
153impl_slice_index!((Bound<usize>, Bound<usize>));
154impl_slice_index!(Range<usize>);
155impl_slice_index!(RangeFrom<usize>);
156impl_slice_index!(RangeFull);
157impl_slice_index!(RangeInclusive<usize>);
158impl_slice_index!(RangeTo<usize>);
159impl_slice_index!(RangeToInclusive<usize>);
mdarray/index/slice.rs

mdarray/index/
slice.rs
