soa_rs/
index.rs

1use crate::{Slice, SliceMut, SliceRef, SoaRaw, Soars};
2use std::{
3    marker::PhantomData,
4    ops::{Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive},
5};
6
7/// A helper trait for indexing operations.
8pub trait SoaIndex<T>
9where
10    T: Soars,
11{
12    /// The output type returned by non-`mut` methods.
13    type Output<'a>
14    where
15        T: 'a;
16
17    /// The output type returned by `mut` methods.
18    type OutputMut<'a>
19    where
20        T: 'a;
21
22    /// Returns the output at this location, if in bounds.
23    fn get(self, slice: &Slice<T>) -> Option<Self::Output<'_>>;
24
25    /// Returns the mutable output at this location, if in bounds.
26    fn get_mut(self, slice: &mut Slice<T>) -> Option<Self::OutputMut<'_>>;
27}
28
29impl<T> SoaIndex<T> for usize
30where
31    T: Soars,
32{
33    type Output<'a>
34        = T::Ref<'a>
35    where
36        T: 'a;
37
38    type OutputMut<'a>
39        = T::RefMut<'a>
40    where
41        T: 'a;
42
43    #[inline]
44    fn get(self, slice: &Slice<T>) -> Option<Self::Output<'_>> {
45        if self < slice.len() {
46            // SAFETY: Offsetting by less than the length leaves
47            // at least one item to read
48            Some(unsafe { slice.raw().offset(self).get_ref() })
49        } else {
50            None
51        }
52    }
53
54    #[inline]
55    fn get_mut(self, slice: &mut Slice<T>) -> Option<Self::OutputMut<'_>> {
56        if self < slice.len() {
57            // SAFETY: Offsetting by less than the length leaves
58            // at least one item to read
59            Some(unsafe { slice.raw().offset(self).get_mut() })
60        } else {
61            None
62        }
63    }
64}
65
66impl<T> SoaIndex<T> for RangeFull
67where
68    T: Soars,
69{
70    type Output<'a>
71        = SliceRef<'a, T>
72    where
73        T: 'a;
74
75    type OutputMut<'a>
76        = SliceMut<'a, T>
77    where
78        T: 'a;
79
80    #[inline]
81    fn get(self, slice: &Slice<T>) -> Option<Self::Output<'_>> {
82        Some(SliceRef {
83            // SAFETY: The lifetime is bound to the given slice
84            slice: unsafe { slice.as_sized() },
85            len: slice.len(),
86            marker: PhantomData,
87        })
88    }
89
90    #[inline]
91    fn get_mut(self, slice: &mut Slice<T>) -> Option<Self::OutputMut<'_>> {
92        Some(SliceMut {
93            // SAFETY: The lifetime is bound to the given slice
94            slice: unsafe { slice.as_sized() },
95            len: slice.len(),
96            marker: PhantomData,
97        })
98    }
99}
100
101impl<T> SoaIndex<T> for Range<usize>
102where
103    T: Soars,
104{
105    type Output<'a>
106        = SliceRef<'a, T>
107    where
108        T: 'a;
109
110    type OutputMut<'a>
111        = SliceMut<'a, T>
112    where
113        T: 'a;
114
115    #[inline]
116    fn get(self, slice: &Slice<T>) -> Option<Self::Output<'_>> {
117        let len = self.len();
118        (len + self.start <= slice.len()).then(|| SliceRef {
119            // SAFETY: The above bounds check ensures we won't be able to access
120            // the slice out of bounds.
121            slice: Slice::with_raw(unsafe { slice.raw.offset(self.start) }),
122            len,
123            marker: PhantomData,
124        })
125    }
126
127    #[inline]
128    fn get_mut(self, slice: &mut Slice<T>) -> Option<Self::OutputMut<'_>> {
129        self.get(slice).map(|s| SliceMut {
130            slice: unsafe { s.as_sized() },
131            len: s.len(),
132            marker: PhantomData,
133        })
134    }
135}
136
137impl<T> SoaIndex<T> for RangeTo<usize>
138where
139    T: Soars,
140{
141    type Output<'a>
142        = SliceRef<'a, T>
143    where
144        T: 'a;
145
146    type OutputMut<'a>
147        = SliceMut<'a, T>
148    where
149        T: 'a;
150
151    #[inline]
152    fn get(self, slice: &Slice<T>) -> Option<Self::Output<'_>> {
153        (0..self.end).get(slice)
154    }
155
156    #[inline]
157    fn get_mut(self, slice: &mut Slice<T>) -> Option<Self::OutputMut<'_>> {
158        (0..self.end).get_mut(slice)
159    }
160}
161
162impl<T> SoaIndex<T> for RangeToInclusive<usize>
163where
164    T: Soars,
165{
166    type Output<'a>
167        = SliceRef<'a, T>
168    where
169        T: 'a;
170
171    type OutputMut<'a>
172        = SliceMut<'a, T>
173    where
174        T: 'a;
175
176    #[inline]
177    fn get(self, slice: &Slice<T>) -> Option<Self::Output<'_>> {
178        (0..self.end + 1).get(slice)
179    }
180
181    #[inline]
182    fn get_mut(self, slice: &mut Slice<T>) -> Option<Self::OutputMut<'_>> {
183        (0..self.end + 1).get_mut(slice)
184    }
185}
186
187impl<T> SoaIndex<T> for RangeFrom<usize>
188where
189    T: Soars,
190{
191    type Output<'a>
192        = SliceRef<'a, T>
193    where
194        T: 'a;
195
196    type OutputMut<'a>
197        = SliceMut<'a, T>
198    where
199        T: 'a;
200
201    #[inline]
202    fn get(self, slice: &Slice<T>) -> Option<Self::Output<'_>> {
203        (self.start..slice.len()).get(slice)
204    }
205
206    #[inline]
207    fn get_mut(self, slice: &mut Slice<T>) -> Option<Self::OutputMut<'_>> {
208        (self.start..slice.len()).get_mut(slice)
209    }
210}
211
212impl<T> SoaIndex<T> for RangeInclusive<usize>
213where
214    T: Soars,
215{
216    type Output<'a>
217        = SliceRef<'a, T>
218    where
219        T: 'a;
220
221    type OutputMut<'a>
222        = SliceMut<'a, T>
223    where
224        T: 'a;
225
226    #[inline]
227    fn get(self, slice: &Slice<T>) -> Option<Self::Output<'_>> {
228        (*self.start()..*self.end() + 1).get(slice)
229    }
230
231    #[inline]
232    fn get_mut(self, slice: &mut Slice<T>) -> Option<Self::OutputMut<'_>> {
233        (*self.start()..*self.end() + 1).get_mut(slice)
234    }
235}