cubecl_core/frontend/iter/adapters/
map.rs

1use crate::prelude::*;
2
3/// The expand half of [`Map`](core::iter::Map).
4#[derive(new)]
5pub struct MapExpand<I, F> {
6    iter: I,
7    f: F,
8}
9
10impl<I, F, B> Iterable for MapExpand<I, F>
11where
12    I: Iterable,
13    F: FnMut(&Scope, I::Item) -> B,
14{
15    type Item = B;
16
17    fn expand(mut self, scope: &Scope, body: &mut dyn FnMut(&Scope, B)) {
18        self.iter.expand(scope, &mut |scope, item| {
19            let mapped = (self.f)(scope, item);
20            body(scope, mapped);
21        });
22    }
23
24    fn expand_unroll(mut self, scope: &Scope, body: &mut dyn FnMut(&Scope, B)) {
25        self.iter.expand_unroll(scope, &mut |scope, item| {
26            let mapped = (self.f)(scope, item);
27            body(scope, mapped);
28        });
29    }
30}
31
32impl<I, F, B> IteratorExpand for MapExpand<I, F>
33where
34    I: IteratorExpand,
35    F: FnMut(&Scope, I::Item) -> B,
36{
37}
38
39impl<I, F, B> IndexedIteratorExpand for MapExpand<I, F>
40where
41    I: IndexedIteratorExpand,
42    F: FnMut(&Scope, I::Item) -> B,
43{
44    fn __expand_len_method(&self, scope: &Scope) -> NativeExpand<usize> {
45        self.iter.__expand_len_method(scope)
46    }
47
48    fn get(&mut self, scope: &Scope, index: NativeExpand<usize>) -> B {
49        let item = self.iter.get(scope, index);
50        (self.f)(scope, item)
51    }
52}
cubecl_core/frontend/iter/adapters/map.rs

cubecl_core/frontend/iter/adapters/
map.rs
