Skip to main content

frozone_derive/
lib.rs

1#![no_std]
2extern crate proc_macro2;
3use proc_macro::TokenStream;
4use quote::quote;
5
6#[proc_macro_derive(Freezable, attributes(assume_frozen))]
7pub fn derive_freezable(input: TokenStream) -> TokenStream {
8    let ast: syn::DeriveInput = syn::parse_macro_input!(input);
9    let name = &ast.ident;
10    let generics = ast.generics.split_for_impl();
11
12    match ast.data {
13        syn::Data::Struct(data) => derive_freezable_struct(data, name, generics),
14        syn::Data::Enum(data) => derive_freezable_enum(data, name, generics),
15        _ => unimplemented!(),
16    }
17}
18
19fn derive_freezable_enum(
20    data: syn::DataEnum,
21    name: &syn::Ident,
22    generics: (
23        syn::ImplGenerics,
24        syn::TypeGenerics,
25        Option<&syn::WhereClause>,
26    ),
27) -> TokenStream {
28    let variants_names_and_freezes = data.variants.iter().map(|f| {
29        let name = &f.ident;
30        if f.attrs.iter().any(|a| a.path().is_ident("assume_frozen")) {
31            quote! {
32                (stringify!(#name), 0)
33            }
34        } else {
35            let discriminant = f
36                .discriminant
37                .as_ref()
38                .map(|eq_d| eq_d.1.clone())
39                .map(|d| {
40                    quote! {
41                    use core::hash::Hasher;
42                    let mut hasher = core::hash::SipHasher::new();
43                        (#d).hash(&mut hasher);
44                    hasher.finish()
45                    }
46                })
47                .unwrap_or(quote! {0});
48            let variant_fields = f.fields.iter().map(|g| {
49                let g_ty = &g.ty;
50                quote! {
51                    <#g_ty as frozone::Freezable>::freeze()
52                }
53            });
54
55            quote! {
56                (stringify!(#name), {
57                    let mut hasher = core::hash::SipHasher::new();
58
59                    #discriminant.hash(&mut hasher);
60                    [#(#variant_fields,)*].iter().for_each(|x: &u64| {
61                        x.hash(&mut hasher);
62                    });
63                    hasher.finish()
64                })
65            }
66        }
67    });
68
69    // #[cfg(feature = "test")]
70    // for v in variants_names_and_freezes.clone() {
71    //     println!("here");
72    //     print!("variants => {}", pretty_print(&v));
73    //     println!("\nhere - end");
74    // }
75    let (impl_generics, type_generics, where_clause) = generics;
76    let generated = quote! {
77        impl #impl_generics frozone::Freezable for #name #type_generics #where_clause {
78            fn freeze() -> u64 {
79                use core::hash::{Hash, Hasher};
80
81                [#(#variants_names_and_freezes,)*].iter().fold(0u64, |acc, x| {
82                    let mut hasher = core::hash::SipHasher::new();
83                    x.0.hash(&mut hasher);
84                    x.1.hash(&mut hasher);
85                    acc.overflowing_add(hasher.finish()).0
86                })
87            }
88        }
89    };
90    let g: proc_macro2::TokenStream = generated.into();
91    // #[cfg(test)]
92    // {
93    // #[cfg(feature = "test")]
94    // print!("AST => {}", pretty_print(&g));
95    g.into()
96    // }
97    // #[cfg(not(test))]
98    // g
99}
100
101fn derive_freezable_struct(
102    data: syn::DataStruct,
103    name: &syn::Ident,
104    generics: (
105        syn::ImplGenerics,
106        syn::TypeGenerics,
107        Option<&syn::WhereClause>,
108    ),
109) -> TokenStream {
110    let fields = data.fields.iter().map(|f| {
111        let name = &f.ident;
112        let ty = &f.ty;
113        if f.attrs.iter().any(|a| a.path().is_ident("assume_frozen")) {
114            quote! {
115                (stringify!(#name), 0)
116            }
117        } else {
118            quote! {
119                (stringify!(#name), <#ty as frozone::Freezable>::freeze())
120            }
121        }
122    });
123
124    let (impl_generics, type_generics, where_clause) = generics;
125    let generated = quote! {
126        impl #impl_generics frozone::Freezable for #name #type_generics #where_clause {
127            fn freeze() -> u64 {
128                use core::hash::{Hash, Hasher};
129
130                // stringify!( [#(#fields,)*]);
131                [#(#fields,)*].iter().fold(0u64, |acc, x| {
132                    let mut hasher = core::hash::SipHasher::new();
133                    x.0.hash(&mut hasher);
134                    x.1.hash(&mut hasher);
135                    acc.overflowing_add(hasher.finish()).0
136                })
137            }
138        }
139    };
140
141    generated.into()
142}
143
144// #[cfg(feature = "test")]
145// fn pretty_print(ts: &proc_macro2::TokenStream) -> String {
146//     let file = match syn::parse_file(&ts.to_string()) {
147//         Ok(f) => f,
148//         Err(e) => return format!("error parsing tokenstream: {e:?}"),
149//     };
150//     prettyplease::unparse(&file)
151// }