tea_core/backends_impl/
vec.rs

1use std::mem::MaybeUninit;
2
3use crate::prelude::*;
4
5macro_rules! impl_vec1 {
6    (to_iter $($ty: ty),* $(,)?) => {
7        $(
8            impl<T> GetLen for $ty {
9                #[inline]
10                fn len(&self) -> usize {
11                    (*self).len()
12                }
13            }
14
15            impl<T: Clone> TIter<T> for $ty {
16                #[inline]
17                fn titer(&self) -> impl TIterator<Item = T>
18                {
19                    self.iter().cloned()
20                }
21            }
22        )*
23    };
24
25    (view $($({$N: ident})? $(--$slice: ident)? $ty: ty),* $(,)?) => {
26        $(
27            impl<T: Clone $(, const $N: usize)?> Vec1View<T> for $ty {
28                type SliceOutput<'a> = &'a <Self as std::ops::Index<std::ops::Range<usize>>>::Output where Self: 'a;
29
30                #[inline]
31                fn slice<'a>(&'a self, start: usize, end: usize) -> TResult<Self::SliceOutput<'a>>
32                where
33                    T: 'a,
34                {
35                    use std::ops::Index;
36                    Ok(self.index(start..end))
37                }
38
39                #[inline]
40                unsafe fn uslice<'a>(&'a self, start: usize, end: usize) -> TResult<Self::SliceOutput<'a>>
41                where
42                    T: 'a,
43                { unsafe {
44                    Ok(self.get_unchecked(start..end))
45                }}
46
47                #[inline]
48                fn get_backend_name(&self) -> &'static str {
49                    "vec"
50                }
51
52                #[inline]
53                unsafe fn uget(&self, index: usize) -> T { unsafe {
54                    self.get_unchecked(index).clone()
55                }}
56
57                $(#[inline]
58                fn $slice(&self) -> Option<&[T]> {
59                    Some(self)
60                })?
61
62
63                // this should be a faster implemention than default as
64                // we read value directly by ptr
65                #[inline]
66                fn rolling_custom<'a, O: Vec1<OT>, OT: Clone, F>(
67                    &'a self,
68                    window: usize,
69                    f: F,
70                    out: Option<O::UninitRefMut<'_>>,
71                ) -> Option<O>
72                where
73                    F: FnMut(Self::SliceOutput<'a>) -> OT,
74                    T: 'a,
75                {
76                    let len = self.len();
77                    if let Some(out) = out {
78                        self.rolling_custom_to::<O, _, _>(window, f, out);
79                        None
80                    } else {
81                        use crate::prelude::UninitVec;
82                        let mut out = O::uninit(len);
83                        self.rolling_custom_to::<O, _, _>(window, f, O::uninit_ref_mut(&mut out));
84                        Some(unsafe { out.assume_init() })
85                    }
86                }
87
88                #[inline]
89                // this should be a faster implemention than default as
90                // we read value directly by ptr
91                fn rolling_apply<O: Vec1<OT>, OT, F>(
92                    &self,
93                    window: usize,
94                    f: F,
95                    out: Option<O::UninitRefMut<'_>>,
96                ) -> Option<O>
97                where
98                    F: FnMut(Option<T>, T) -> OT,
99                {
100                    let len = self.len();
101                    if let Some(out) = out {
102                        self.rolling_apply_to::<O, _, _>(window, f, out);
103                        None
104                    } else {
105                        let mut out = O::uninit(len);
106                        self.rolling_apply_to::<O, _, _>(window, f, O::uninit_ref_mut(&mut out));
107                        Some(unsafe { out.assume_init() })
108                    }
109                }
110
111                #[inline]
112                // this should be a faster implemention than default as
113                // we read value directly by ptr
114                fn rolling2_apply<O: Vec1<OT>, OT, V2: Vec1View<T2>, T2, F>(
115                    &self,
116                    other: &V2,
117                    window: usize,
118                    f: F,
119                    out: Option<O::UninitRefMut<'_>>,
120                ) -> Option<O>
121                where
122                    F: FnMut(Option<(T, T2)>, (T, T2)) -> OT,
123                {
124                    let len = self.len();
125                    if let Some(out) = out {
126                        self.rolling2_apply_to::<O, _, _, _, _>(other, window, f, out);
127                        None
128                    } else {
129                        let mut out = O::uninit(len);
130                        self.rolling2_apply_to::<O, _, _, _, _>(other, window, f, O::uninit_ref_mut(&mut out));
131                        Some(unsafe { out.assume_init() })
132                    }
133                }
134
135
136                // this should be a faster implemention than default as
137                // we read value directly by ptr
138                #[inline]
139                fn rolling_apply_idx<O: Vec1<OT>, OT, F>(
140                    &self,
141                    window: usize,
142                    f: F,
143                    out: Option<O::UninitRefMut<'_>>,
144                ) -> Option<O>
145                where
146                    F: FnMut(Option<usize>, usize, T) -> OT,
147                {
148                    let len = self.len();
149                    if let Some(out) = out {
150                        self.rolling_apply_idx_to::<O, _, _>(window, f, out);
151                        None
152                    } else {
153                        let mut out = O::uninit(len);
154                        self.rolling_apply_idx_to::<O, _, _>(window, f, O::uninit_ref_mut(&mut out));
155                        Some(unsafe { out.assume_init() })
156                    }
157                }
158
159
160                // this should be a faster implemention than default as
161                // we read value directly by ptr
162                #[inline]
163                fn rolling2_apply_idx<O: Vec1<OT>, OT, V2: Vec1View<T2>, T2, F>(
164                    &self,
165                    other: &V2,
166                    window: usize,
167                    f: F,
168                    out: Option<O::UninitRefMut<'_>>,
169                ) -> Option<O>
170                where
171                    F: FnMut(Option<usize>, usize, (T, T2)) -> OT,
172                {
173                    let len = self.len();
174                    if let Some(out) = out {
175                        self.rolling2_apply_idx_to::<O, _, _, _, _>(other, window, f, out);
176                        None
177                    } else {
178                        let mut out = O::uninit(len);
179                        self.rolling2_apply_idx_to::<O, _, _, _, _>(other, window, f, O::uninit_ref_mut(&mut out));
180                        Some(unsafe { out.assume_init() })
181                    }
182                }
183            }
184        )*
185    };
186}
187
188impl<T, const N: usize> GetLen for [T; N] {
189    #[inline]
190    fn len(&self) -> usize {
191        N
192    }
193}
194
195impl<T: Clone, const N: usize> TIter<T> for [T; N] {
196    #[inline]
197    fn titer(&self) -> impl TIterator<Item = T> {
198        self.iter().cloned()
199    }
200}
201
202impl_vec1!(
203    to_iter
204    Vec<T>,
205    [T],
206    &[T],
207);
208
209impl<T> GetLen for &mut [T] {
210    #[inline]
211    fn len(&self) -> usize {
212        (**self).len()
213    }
214}
215
216impl<T: Clone> TIter<T> for &mut [T] {
217    #[inline]
218    fn titer(&self) -> impl TIterator<Item = T> {
219        self.iter().cloned()
220    }
221}
222
223impl_vec1!(
224    view
225    --try_as_slice Vec<T>,
226    --try_as_slice [T],
227    // --try_as_slice &[T],
228    {N} [T; N]
229);
230
231impl<'a, T: Clone + 'a> Vec1Mut<'a, T> for Vec<T> {
232    #[inline]
233    unsafe fn uget_mut(&mut self, index: usize) -> &mut T {
234        unsafe { self.get_unchecked_mut(index) }
235    }
236
237    #[inline]
238    fn try_as_slice_mut(&mut self) -> Option<&mut [T]> {
239        Some(self.as_mut_slice())
240    }
241}
242
243impl<T: Clone> Vec1<T> for Vec<T> {
244    type Uninit = Vec<MaybeUninit<T>>;
245    type UninitRefMut<'a>
246        = &'a mut [MaybeUninit<T>]
247    where
248        T: 'a;
249
250    #[inline]
251    fn collect_from_iter<I: Iterator<Item = T>>(iter: I) -> Self {
252        iter.collect()
253    }
254
255    #[inline]
256    fn try_collect_from_iter<I: Iterator<Item = TResult<T>>>(iter: I) -> TResult<Self> {
257        iter.collect()
258    }
259
260    #[inline]
261    fn uninit(len: usize) -> Self::Uninit {
262        let mut v = Vec::with_capacity(len);
263        unsafe {
264            v.set_len(len);
265        }
266        v
267    }
268
269    #[inline]
270    fn uninit_ref_mut(uninit_vec: &mut Self::Uninit) -> Self::UninitRefMut<'_> {
271        uninit_vec
272    }
273
274    #[inline]
275    fn collect_from_trusted<I: Iterator<Item = T> + TrustedLen>(iter: I) -> Self {
276        iter.collect_trusted_to_vec()
277    }
278
279    #[inline]
280    fn try_collect_from_trusted<I: Iterator<Item = TResult<T>> + TrustedLen>(
281        iter: I,
282    ) -> TResult<Self> {
283        iter.try_collect_trusted_to_vec()
284    }
285
286    #[inline]
287    fn empty() -> Self {
288        Vec::new()
289    }
290}
291
292impl<T: Clone> UninitVec<T> for Vec<MaybeUninit<T>> {
293    type Vec = Vec<T>;
294
295    #[inline]
296    unsafe fn assume_init(self) -> Self::Vec {
297        unsafe { std::mem::transmute::<Vec<MaybeUninit<T>>, Vec<T>>(self) }
298    }
299
300    #[inline]
301    unsafe fn uset(&mut self, idx: usize, v: T) {
302        unsafe {
303            let ele = self.get_unchecked_mut(idx);
304            ele.write(v);
305        }
306    }
307}
308
309impl<T> UninitRefMut<T> for &mut [MaybeUninit<T>] {
310    #[inline]
311    unsafe fn uset(&mut self, idx: usize, v: T) {
312        unsafe {
313            let ele = self.get_unchecked_mut(idx);
314            ele.write(v);
315        }
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use crate::prelude::*;
322
323    #[test]
324    fn test_get() {
325        let mut data = vec![1, 2, 3, 4, 5];
326        {
327            let view_mut: &mut [_] = &mut data;
328            assert_eq!(GetLen::len(view_mut), 5);
329            assert_eq!(Vec1View::get(view_mut, 2).unwrap(), 3);
330        }
331        let view = &data;
332        assert_eq!(GetLen::len(&data), 5);
333        assert_eq!(GetLen::len(&[1, 2]), 2);
334        assert_eq!(view.get(0).unwrap(), 1);
335        assert_eq!([1, 2].get(0).unwrap(), 1);
336        let slice = view.as_slice();
337        assert_eq!(unsafe { slice.uget(2) }, 3);
338    }
339
340    #[test]
341    fn test_slice() {
342        let v = vec![1, 2, 4, 5, 2];
343        let res = v.slice(0, 3).unwrap().titer().collect_trusted_to_vec();
344        assert_eq!(&res, &[1, 2, 4]);
345        let res = v.slice(2, 4).unwrap().titer().collect_trusted_to_vec();
346        assert_eq!(&res, &[4, 5]);
347    }
348
349    #[test]
350    fn test_get_mut() {
351        let mut data = vec![1, 2, 3];
352        *unsafe { Vec1Mut::uget_mut(&mut data, 1) } = 4;
353        assert_eq!(data[1], 4);
354        let mut_ref = &mut data;
355        *unsafe { Vec1Mut::uget_mut(mut_ref, 1) } = 4;
356    }
357
358    #[test]
359    fn test_collect() {
360        let data = (0..5).collect_vec1::<Vec<_>>();
361        assert_eq!(data, vec![0, 1, 2, 3, 4]);
362        let data = (0..5).collect_trusted_vec1::<Vec<_>>();
363        assert_eq!(data, vec![0, 1, 2, 3, 4]);
364        let v: Vec<i32> = vec![];
365        let data: Vec<i32> = Vec::empty();
366        assert_eq!(data, v);
367        let data = vec![Some(1.), None, Some(2.)].collect_vec1_opt::<Vec<f64>>();
368        assert!(data[1].is_nan());
369        assert_eq!(data[2], 2.);
370        let data = vec![Ok(1), Ok(2), Err(terr!("err")), Ok(3)];
371        let v: TResult<Vec<_>> = data.try_collect_vec1();
372        assert!(v.is_err());
373        let data = vec![Ok(1), Ok(2), Ok(3)];
374        let v: TResult<Vec<_>> = data.try_collect_trusted_vec1();
375        assert_eq!(v.unwrap(), vec![1, 2, 3]);
376    }
377
378    #[test]
379    fn test_iter_cast() {
380        let data = vec![1, 2, 3, 4, 5];
381        let out: Vec<_> = data.iter_cast::<f64>().collect_trusted_vec1();
382        assert_eq!(out, vec![1., 2., 3., 4., 5.]);
383        let data = vec![Some(1), Some(2), None];
384        let out: Vec<_> = data.opt_iter_cast::<f64>().collect_vec1();
385        assert_eq!(out, vec![Some(1.), Some(2.), None])
386    }
387
388    #[test]
389    fn test_uninit() {
390        let mut data = Vec::<i32>::uninit(2);
391        unsafe { data.uset(0, 1) };
392        unsafe { data.uset(1, 2) };
393        let data: Vec<_> = unsafe { data.assume_init() };
394        assert_eq!(data, vec![1, 2]);
395    }
396
397    #[test]
398    fn test_rolling_custom() {
399        let data = vec![1, 2, 3, 4, 5];
400        let out: Vec<_> = data
401            .rolling_custom(3, |s| s.titer().vsum().unwrap(), None)
402            .unwrap();
403        assert_eq!(out, vec![1, 3, 6, 9, 12]);
404    }
405}