cubecl_core/frontend/iter/traits/
iterator.rs

1use crate::frontend::iter::{CopiedExpand, FilterExpand, MapExpand};
2use crate::frontend::{break_expand, if_expand};
3use crate::prelude::*;
4
5/// An argument of a generic iterator method.
6///
7/// `#[cube]` passes every argument other than a closure as `(arg).into()`, which can't infer a bare
8/// generic parameter. The single `From` impl of this wrapper leaves the compiler with one candidate,
9/// which infers the parameter from the argument.
10pub struct IterArg<T>(pub(super) T);
11
12impl<T> From<T> for IterArg<T> {
13    fn from(value: T) -> Self {
14        Self(value)
15    }
16}
17
18/// The expand half of [`Iterator`](core::iter::Iterator). Consumers expand a single loop through
19/// [`Iterable::expand`].
20pub trait IteratorExpand: Iterable {
21    /// See [`Iterator::map`](core::iter::Iterator::map).
22    fn __expand_map_method<B, F>(self, _scope: &Scope, f: F) -> MapExpand<Self, F>
23    where
24        F: FnMut(&Scope, Self::Item) -> B,
25    {
26        MapExpand::new(self, f)
27    }
28
29    /// See [`Iterator::filter`](core::iter::Iterator::filter).
30    fn __expand_filter_method<P>(self, _scope: &Scope, predicate: P) -> FilterExpand<Self, P>
31    where
32        P: FnMut(&Scope, &Self::Item) -> NativeExpand<bool>,
33    {
34        FilterExpand::new(self, predicate)
35    }
36
37    /// See [`Iterator::copied`](core::iter::Iterator::copied).
38    fn __expand_copied_method<'a, T>(self, _scope: &Scope) -> CopiedExpand<Self>
39    where
40        T: DerefExpand<Target = T> + 'a,
41        Self: Iterable<Item = &'a T>,
42    {
43        CopiedExpand::new(self)
44    }
45
46    /// See [`Iterator::for_each`](core::iter::Iterator::for_each).
47    fn __expand_for_each_method<F>(self, scope: &Scope, mut f: F)
48    where
49        F: FnMut(&Scope, Self::Item),
50    {
51        self.expand(scope, &mut |scope, item| f(scope, item));
52    }
53
54    /// See [`Iterator::fold`](core::iter::Iterator::fold).
55    fn __expand_fold_method<B, F>(self, scope: &Scope, init: IterArg<B>, mut f: F) -> B::Expand
56    where
57        B: IntoExpand<Expand: Assign + IntoMut + Clone>,
58        F: FnMut(&Scope, B::Expand, Self::Item) -> B::Expand,
59    {
60        let mut accum = init.0.into_expand(scope).into_mut(scope);
61        self.expand(scope, &mut |scope, item| {
62            let next = f(scope, accum.clone(), item);
63            accum.__expand_assign_method(scope, next);
64        });
65        accum
66    }
67
68    /// See [`Iterator::sum`](core::iter::Iterator::sum).
69    ///
70    /// The result must be a concrete type such as `sum::<u32>()`, because a generic `F: Float` or
71    /// `I: Int` doesn't implement [`Sum`](core::iter::Sum). Use `fold` for those.
72    fn __expand_sum_method<S>(self, scope: &Scope) -> NativeExpand<S>
73    where
74        S: Numeric,
75        Self::Item: IntoNumericExpand<S>,
76    {
77        let mut accum = S::__expand_from_int(scope, 0i64.into()).into_mut(scope);
78        self.expand(scope, &mut |scope, item| {
79            let value = item.into_numeric_expand(scope);
80            accum.__expand_add_assign_method(scope, value);
81        });
82        accum
83    }
84
85    /// See [`Iterator::product`](core::iter::Iterator::product).
86    ///
87    /// The result must be a concrete type for the same reason as `sum`.
88    fn __expand_product_method<P>(self, scope: &Scope) -> NativeExpand<P>
89    where
90        P: Numeric,
91        Self::Item: IntoNumericExpand<P>,
92    {
93        let mut accum = P::__expand_from_int(scope, 1i64.into()).into_mut(scope);
94        self.expand(scope, &mut |scope, item| {
95            let value = item.into_numeric_expand(scope);
96            accum.__expand_mul_assign_method(scope, value);
97        });
98        accum
99    }
100
101    /// See [`Iterator::count`](core::iter::Iterator::count).
102    fn __expand_count_method(self, scope: &Scope) -> NativeExpand<usize> {
103        let mut count = NativeExpand::<usize>::from_lit(scope, 0).into_mut(scope);
104        self.expand(scope, &mut |scope, _| {
105            let one = NativeExpand::<usize>::from_lit(scope, 1);
106            count.__expand_add_assign_method(scope, one);
107        });
108        count
109    }
110
111    /// See [`Iterator::any`](core::iter::Iterator::any). Stops at the first match.
112    fn __expand_any_method<F>(self, scope: &Scope, mut f: F) -> NativeExpand<bool>
113    where
114        F: FnMut(&Scope, Self::Item) -> NativeExpand<bool>,
115    {
116        let mut found = false.__expand_runtime_method(scope).into_mut(scope);
117        self.expand(scope, &mut |scope, item| {
118            let matched = f(scope, item);
119            if_expand(scope, matched, |scope| {
120                found.__expand_assign_method(scope, true.into());
121                break_expand(scope);
122            });
123        });
124        found
125    }
126
127    /// See [`Iterator::all`](core::iter::Iterator::all). Stops at the first item that doesn't match.
128    fn __expand_all_method<F>(self, scope: &Scope, mut f: F) -> NativeExpand<bool>
129    where
130        F: FnMut(&Scope, Self::Item) -> NativeExpand<bool>,
131    {
132        self.__expand_any_method(scope, |scope, item| {
133            f(scope, item).__expand_not_method(scope)
134        })
135        .__expand_not_method(scope)
136    }
137}
cubecl_core/frontend/iter/traits/iterator.rs

cubecl_core/frontend/iter/traits/
iterator.rs
