cubecl_core/frontend/container/slice/
iter.rs

1use crate::frontend::iter::impl_iterable_indexed;
2use crate::prelude::*;
3
4use super::SliceExpand;
5
6impl<E: CubePrimitive> SliceExpand<E> {
7    /// See [`slice::iter`].
8    pub fn __expand_iter_method(&self, _scope: &Scope) -> IterExpand<'_, E> {
9        IterExpand { slice: self }
10    }
11
12    /// See [`slice::iter_mut`].
13    pub fn __expand_iter_mut_method(&mut self, _scope: &Scope) -> IterMutExpand<'_, E> {
14        IterMutExpand { slice: self }
15    }
16}
17
18/// The expand half of [`Iter`](core::slice::Iter).
19pub struct IterExpand<'a, E: CubePrimitive> {
20    slice: &'a SliceExpand<E>,
21}
22
23impl_iterable_indexed!(['a, E: CubePrimitive] IterExpand<'a, E>, &'a NativeExpand<E>);
24
25impl<E: CubePrimitive> IndexedIteratorExpand for IterExpand<'_, E> {
26    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
27        self.slice.__expand_len_method(scope)
28    }
29
30    fn get(&mut self, scope: &Scope, index: NativeExpand<usize>) -> Self::Item {
31        let slice = self.slice;
32        slice.__expand_index_method(scope, index)
33    }
34}
35
36/// The expand half of [`IterMut`](core::slice::IterMut).
37pub struct IterMutExpand<'a, E: CubePrimitive> {
38    // Shared because `get` hands out many `&'a mut` references, each to a separate kernel value
39    // rather than to anything in this slice expand.
40    slice: &'a SliceExpand<E>,
41}
42
43impl_iterable_indexed!(['a, E: CubePrimitive] IterMutExpand<'a, E>, &'a mut NativeExpand<E>);
44
45impl<E: CubePrimitive> IndexedIteratorExpand for IterMutExpand<'_, E> {
46    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
47        self.slice.__expand_len_method(scope)
48    }
49
50    fn get(&mut self, scope: &Scope, index: NativeExpand<usize>) -> Self::Item {
51        let slice = self.slice;
52        slice
53            .__expand_as_mut_unchecked_method(scope)
54            .__expand_index_mut_method(scope, index)
55    }
56}
57
58impl<'a, E: CubePrimitive> IntoIteratorExpand for &'a SliceExpand<E> {
59    type IntoIter = IterExpand<'a, E>;
60
61    fn __expand_into_iter_method(self, scope: &Scope) -> Self::IntoIter {
62        self.__expand_iter_method(scope)
63    }
64}
65
66impl<'a, E: CubePrimitive> IntoIteratorExpand for &'a mut SliceExpand<E> {
67    type IntoIter = IterMutExpand<'a, E>;
68
69    fn __expand_into_iter_method(self, scope: &Scope) -> Self::IntoIter {
70        self.__expand_iter_mut_method(scope)
71    }
72}
cubecl_core/frontend/container/slice/iter.rs

cubecl_core/frontend/container/slice/
iter.rs
