cubecl_core/frontend/iter/traits/
indexed.rs

1use crate::frontend::iter::{EnumerateExpand, RevExpand, SkipExpand, TakeExpand, ZipExpand};
2use crate::prelude::*;
3
4/// An iterator with a known length and random access to its items.
5///
6/// Slice iterators and ranges are indexed, and adapters over them stay indexed. This has no
7/// counterpart in the public API of [`core::iter`], which `next` makes possible for any iterator.
8pub trait IndexedIteratorExpand: IteratorExpand {
9    /// See [`ExactSizeIterator::len`](core::iter::ExactSizeIterator::len).
10    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize>;
11
12    /// The item at `index`, which must be less than the length.
13    fn get(&mut self, scope: &Scope, index: NativeExpand<usize>) -> Self::Item;
14
15    /// See [`Iterator::zip`](core::iter::Iterator::zip).
16    fn __expand_zip_method<U>(
17        self,
18        scope: &Scope,
19        other: IterArg<U>,
20    ) -> ZipExpand<Self, U::IntoIter>
21    where
22        U: IntoIteratorExpand,
23    {
24        ZipExpand::new(self, other.0.__expand_into_iter_method(scope))
25    }
26
27    /// See [`Iterator::enumerate`](core::iter::Iterator::enumerate).
28    fn __expand_enumerate_method(self, _scope: &Scope) -> EnumerateExpand<Self> {
29        EnumerateExpand::new(self)
30    }
31
32    /// See [`Iterator::rev`](core::iter::Iterator::rev).
33    fn __expand_rev_method(self, _scope: &Scope) -> RevExpand<Self> {
34        RevExpand::new(self)
35    }
36
37    /// See [`Iterator::take`](core::iter::Iterator::take).
38    fn __expand_take_method(self, _scope: &Scope, n: NativeExpand<usize>) -> TakeExpand<Self> {
39        TakeExpand::new(self, n)
40    }
41
42    /// See [`Iterator::skip`](core::iter::Iterator::skip).
43    fn __expand_skip_method(self, _scope: &Scope, n: NativeExpand<usize>) -> SkipExpand<Self> {
44        SkipExpand::new(self, n)
45    }
46}
47
48/// Expand the loop of an indexed iterator.
49///
50/// This calls [`Iterable::expand`] rather than [`for_expand`], which leaves out the loop of a
51/// constant length of one that `any` and `all` need to break out of.
52pub(crate) fn expand_indexed<I: IndexedIteratorExpand>(
53    mut iter: I,
54    scope: &Scope,
55    body: &mut dyn FnMut(&Scope, I::Item),
56) {
57    let start = NativeExpand::<usize>::from_lit(scope, 0);
58    let end = iter.__expand_len_method(scope);
59    RangeExpand::new(start, end).expand(scope, &mut |scope, index| {
60        let item = iter.get(scope, index);
61        body(scope, item);
62    });
63}
cubecl_core/frontend/iter/traits/indexed.rs

cubecl_core/frontend/iter/traits/
indexed.rs
