cubecl-macros-internal 0.11.0

Internal procedural macros for CubeCL
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

use syn::{
    ExprClosure, Ident, Token, Type,
    parse::{Parse, ParseStream},
    punctuated::Punctuated,
    token,
};

pub struct ConstEval {
    pub op_type: Type,
    _comma: Token![,],
    _brace: token::Brace,
    pub arms: Punctuated<FoldArm, Token![,]>,
}

pub enum FoldArm {
    ConstEval(ConstEvalArm),
    Custom(CustomArm),
}

pub struct ConstEvalArm {
    pub attr_types: TypeList,
    _colon: Token![:],
    pub closure: ExprClosure,
}

pub struct CustomArm {
    _custom: Ident,
    _colon: Token![:],
    pub closure: ExprClosure,
}

pub struct TypeList(pub Punctuated<AttrType, Token![,]>);

pub struct AttrType {
    pub attr_ty: Type,
    pub subtypes: Punctuated<Type, Token![,]>,
}

impl Parse for ConstEval {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let content;
        Ok(ConstEval {
            op_type: input.parse()?,
            _comma: input.parse()?,
            _brace: syn::braced!(content in input),
            arms: content.parse_terminated(FoldArm::parse, Token![,])?,
        })
    }
}

impl Parse for FoldArm {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let fork = input.fork();
        if let Ok(ident) = fork.parse::<Ident>()
            && ident == "custom"
        {
            Ok(FoldArm::Custom(CustomArm::parse(input)?))
        } else {
            Ok(FoldArm::ConstEval(ConstEvalArm::parse(input)?))
        }
    }
}

impl Parse for ConstEvalArm {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        Ok(ConstEvalArm {
            attr_types: input.parse()?,
            _colon: input.parse()?,
            closure: input.parse()?,
        })
    }
}

impl Parse for CustomArm {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        Ok(CustomArm {
            _custom: input.parse()?,
            _colon: input.parse()?,
            closure: input.parse()?,
        })
    }
}

impl Parse for TypeList {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let types = if input.peek(token::Bracket) {
            let content;
            syn::bracketed!(content in input);
            content.parse_terminated(AttrType::parse, Token![,])?
        } else {
            let mut types = Punctuated::new();
            types.push(input.parse()?);
            types
        };
        Ok(Self(types))
    }
}

impl Parse for AttrType {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let attr_ty: Type = input.parse()?;
        let subtypes = if input.peek(token::Paren) {
            let content;
            syn::parenthesized!(content in input);
            content.parse_terminated(Type::parse, Token![,])?
        } else {
            Punctuated::new()
        };
        Ok(Self { attr_ty, subtypes })
    }
}