cubecl-macros-internal 0.11.3

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

use darling::{
    FromDeriveInput, FromField, FromVariant,
    ast::{Data, Fields},
    util::Flag,
};
use proc_macro2::TokenStream;
use quote::{ToTokens, quote};
use syn::{DeriveInput, Expr, Generics, Ident, Path, Type};

#[derive(FromDeriveInput)]
#[darling(attributes(type_hash))]
pub struct TypeHash {
    pub ident: Ident,
    pub generics: Generics,
    pub data: Data<TypeHashVariant, TypeHashField>,
}

#[derive(FromVariant)]
#[darling(attributes(type_hash))]
pub struct TypeHashVariant {
    pub discriminant: Option<Expr>,
    pub fields: Fields<TypeHashField>,
}

#[derive(FromField)]
#[darling(attributes(type_hash))]
pub struct TypeHashField {
    pub ident: Option<Ident>,
    pub ty: Type,

    pub r#as: Option<Path>,
    pub skip: Flag,
    pub foreign_type: Flag,
}

pub fn type_hash_impl(input: DeriveInput) -> TokenStream {
    let type_hash = TypeHash::from_derive_input(&input).unwrap();
    match type_hash.data {
        Data::Enum(data) => type_hash_enum(&type_hash.ident, &type_hash.generics, &data),
        Data::Struct(fields) => type_hash_struct(&type_hash.ident, &type_hash.generics, &fields),
    }
}

fn type_hash_struct(
    ident: &Ident,
    generics: &Generics,
    data: &Fields<TypeHashField>,
) -> TokenStream {
    let name = ident.to_string();
    let fields = data.iter().flat_map(write_field_hash);
    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
    quote! {
        impl #impl_generics crate::TypeHash for #ident #ty_generics #where_clause {
            fn write_hash(hasher: &mut impl core::hash::Hasher) {
                hasher.write(#name.as_bytes());
                #(#fields)*
            }
        }
    }
}

fn type_hash_enum(ident: &Ident, generics: &Generics, data: &[TypeHashVariant]) -> TokenStream {
    let name = ident.to_string();
    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
    let variants = data.iter().flat_map(|v| {
        v.discriminant
            .iter()
            .map(|discriminant| quote! [core::hash::Hash::hash(&(#discriminant as isize), hasher);])
            .chain(v.fields.iter().map(write_field_hash))
    });
    quote! {
        impl #impl_generics crate::TypeHash for #ident #ty_generics #where_clause{
            fn write_hash(hasher: &mut impl core::hash::Hasher) {
                hasher.write(#name.as_bytes());
                #(#variants)*
            }
        }
    }
}

fn write_field_hash(field: &TypeHashField) -> TokenStream {
    let field_name = field
        .ident
        .as_ref()
        .map(|name| {
            let name = name.to_token_stream().to_string();
            quote! { hasher.write(#name.as_bytes()); }
        })
        .unwrap_or_default();
    if let Some(ty) = &field.r#as {
        quote! {
            #field_name
            <#ty as crate::TypeHash>::write_hash(hasher);
        }
    } else if field.skip.is_present() {
        TokenStream::new()
    } else if field.foreign_type.is_present() {
        let type_str = field.ty.to_token_stream().to_string();
        quote! {
            #field_name
            hasher.write(#type_str.as_bytes());
        }
    } else {
        let field_type = &field.ty;
        quote! {
            #field_name
            <#field_type as crate::TypeHash>::write_hash(hasher);
        }
    }
}