trezoa_frozen_abi_macro/
lib.rs

1extern crate proc_macro;
2
3// This file littered with these essential cfgs so ensure them.
4#[cfg(not(any(RUSTC_WITH_SPECIALIZATION, RUSTC_WITHOUT_SPECIALIZATION)))]
5compile_error!("rustc_version is missing in build dependency and build.rs is not specified");
6
7#[cfg(any(RUSTC_WITH_SPECIALIZATION, RUSTC_WITHOUT_SPECIALIZATION))]
8use proc_macro::TokenStream;
9
10// Define dummy macro_attribute and macro_derive for stable rustc
11
12#[cfg(RUSTC_WITHOUT_SPECIALIZATION)]
13#[proc_macro_attribute]
14pub fn frozen_abi(_attrs: TokenStream, item: TokenStream) -> TokenStream {
15    item
16}
17
18#[cfg(RUSTC_WITHOUT_SPECIALIZATION)]
19#[proc_macro_derive(AbiExample)]
20pub fn derive_abi_sample(_item: TokenStream) -> TokenStream {
21    "".parse().unwrap()
22}
23
24#[cfg(RUSTC_WITHOUT_SPECIALIZATION)]
25#[proc_macro_derive(AbiEnumVisitor)]
26pub fn derive_abi_enum_visitor(_item: TokenStream) -> TokenStream {
27    "".parse().unwrap()
28}
29
30#[cfg(RUSTC_WITH_SPECIALIZATION)]
31use proc_macro2::{Span, TokenStream as TokenStream2, TokenTree};
32#[cfg(RUSTC_WITH_SPECIALIZATION)]
33use quote::{quote, ToTokens};
34#[cfg(RUSTC_WITH_SPECIALIZATION)]
35use syn::{
36    parse_macro_input, Attribute, Error, Fields, Ident, Item, ItemEnum, ItemStruct, ItemType,
37    LitStr, Variant,
38};
39
40#[cfg(RUSTC_WITH_SPECIALIZATION)]
41fn filter_serde_attrs(attrs: &[Attribute]) -> bool {
42    fn contains_skip(tokens: TokenStream2) -> bool {
43        for token in tokens.into_iter() {
44            match token {
45                TokenTree::Group(group) => {
46                    if contains_skip(group.stream()) {
47                        return true;
48                    }
49                }
50                TokenTree::Ident(ident) => {
51                    if ident == "skip" {
52                        return true;
53                    }
54                }
55                TokenTree::Punct(_) | TokenTree::Literal(_) => (),
56            }
57        }
58
59        false
60    }
61
62    for attr in attrs {
63        if !attr.path().is_ident("serde") {
64            continue;
65        }
66
67        if contains_skip(attr.to_token_stream()) {
68            return true;
69        }
70    }
71
72    false
73}
74
75#[cfg(RUSTC_WITH_SPECIALIZATION)]
76fn filter_allow_attrs(attrs: &mut Vec<Attribute>) {
77    attrs.retain(|attr| {
78        let ss = &attr.path().segments.first().unwrap().ident.to_string();
79        ss.starts_with("allow")
80    });
81}
82
83#[cfg(RUSTC_WITH_SPECIALIZATION)]
84fn derive_abi_sample_enum_type(input: ItemEnum) -> TokenStream {
85    let type_name = &input.ident;
86
87    let mut sample_variant = quote! {};
88    let mut sample_variant_found = false;
89
90    for variant in &input.variants {
91        let variant_name = &variant.ident;
92        let variant = &variant.fields;
93        if *variant == Fields::Unit {
94            sample_variant.extend(quote! {
95                #type_name::#variant_name
96            });
97        } else if let Fields::Unnamed(variant_fields) = variant {
98            let mut fields = quote! {};
99            for field in &variant_fields.unnamed {
100                if !(field.ident.is_none() && field.colon_token.is_none()) {
101                    unimplemented!("tuple enum: {:?}", field);
102                }
103                let field_type = &field.ty;
104                fields.extend(quote! {
105                    <#field_type>::example(),
106                });
107            }
108            sample_variant.extend(quote! {
109                #type_name::#variant_name(#fields)
110            });
111        } else if let Fields::Named(variant_fields) = variant {
112            let mut fields = quote! {};
113            for field in &variant_fields.named {
114                if field.ident.is_none() || field.colon_token.is_none() {
115                    unimplemented!("tuple enum: {:?}", field);
116                }
117                let field_type = &field.ty;
118                let field_name = &field.ident;
119                fields.extend(quote! {
120                    #field_name: <#field_type>::example(),
121                });
122            }
123            sample_variant.extend(quote! {
124                #type_name::#variant_name{#fields}
125            });
126        } else {
127            unimplemented!("{:?}", variant);
128        }
129
130        if !sample_variant_found {
131            sample_variant_found = true;
132            break;
133        }
134    }
135
136    if !sample_variant_found {
137        unimplemented!("empty enum");
138    }
139
140    let mut attrs = input.attrs.clone();
141    filter_allow_attrs(&mut attrs);
142    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
143
144    let result = quote! {
145        #[automatically_derived]
146        #( #attrs )*
147        impl #impl_generics ::trezoa_frozen_abi::abi_example::AbiExample for #type_name #ty_generics #where_clause {
148            fn example() -> Self {
149                ::log::info!(
150                    "AbiExample for enum: {}",
151                    std::any::type_name::<#type_name #ty_generics>()
152                );
153                #sample_variant
154            }
155        }
156    };
157    result.into()
158}
159
160#[cfg(RUSTC_WITH_SPECIALIZATION)]
161fn derive_abi_sample_struct_type(input: ItemStruct) -> TokenStream {
162    let type_name = &input.ident;
163    let mut sample_fields = quote! {};
164    let fields = &input.fields;
165
166    match fields {
167        Fields::Named(_) => {
168            for field in fields {
169                let field_name = &field.ident;
170                sample_fields.extend(quote! {
171                    #field_name: AbiExample::example(),
172                });
173            }
174            sample_fields = quote! {
175                { #sample_fields }
176            }
177        }
178        Fields::Unnamed(_) => {
179            for _ in fields {
180                sample_fields.extend(quote! {
181                    AbiExample::example(),
182                });
183            }
184            sample_fields = quote! {
185                ( #sample_fields )
186            }
187        }
188        _ => unimplemented!("fields: {:?}", fields),
189    }
190
191    let mut attrs = input.attrs.clone();
192    filter_allow_attrs(&mut attrs);
193    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
194    let turbofish = ty_generics.as_turbofish();
195
196    let result = quote! {
197        #[automatically_derived]
198        #( #attrs )*
199        impl #impl_generics ::trezoa_frozen_abi::abi_example::AbiExample for #type_name #ty_generics #where_clause {
200            fn example() -> Self {
201                ::log::info!(
202                    "AbiExample for struct: {}",
203                    std::any::type_name::<#type_name #ty_generics>()
204                );
205                use ::trezoa_frozen_abi::abi_example::AbiExample;
206
207                #type_name #turbofish #sample_fields
208            }
209        }
210    };
211
212    result.into()
213}
214
215#[cfg(RUSTC_WITH_SPECIALIZATION)]
216#[proc_macro_derive(AbiExample)]
217pub fn derive_abi_sample(item: TokenStream) -> TokenStream {
218    let item = parse_macro_input!(item as Item);
219
220    match item {
221        Item::Struct(input) => derive_abi_sample_struct_type(input),
222        Item::Enum(input) => derive_abi_sample_enum_type(input),
223        _ => Error::new_spanned(item, "AbiSample isn't applicable; only for struct and enum")
224            .to_compile_error()
225            .into(),
226    }
227}
228
229#[cfg(RUSTC_WITH_SPECIALIZATION)]
230fn do_derive_abi_enum_visitor(input: ItemEnum) -> TokenStream {
231    let type_name = &input.ident;
232    let mut serialized_variants = quote! {};
233    let mut variant_count: u64 = 0;
234    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
235    for variant in &input.variants {
236        // Don't digest a variant with serde(skip)
237        if filter_serde_attrs(&variant.attrs) {
238            continue;
239        };
240        let sample_variant = quote_sample_variant(type_name, &ty_generics, variant);
241        variant_count = if let Some(variant_count) = variant_count.checked_add(1) {
242            variant_count
243        } else {
244            break;
245        };
246        serialized_variants.extend(quote! {
247            #sample_variant;
248            Serialize::serialize(&sample_variant, digester.create_enum_child()?)?;
249        });
250    }
251
252    let type_str = format!("{type_name}");
253    (quote! {
254        impl #impl_generics ::trezoa_frozen_abi::abi_example::AbiEnumVisitor for #type_name #ty_generics #where_clause {
255            fn visit_for_abi(&self, digester: &mut ::trezoa_frozen_abi::abi_digester::AbiDigester) -> ::trezoa_frozen_abi::abi_digester::DigestResult {
256                let enum_name = #type_str;
257                use ::serde::ser::Serialize;
258                use ::trezoa_frozen_abi::abi_example::AbiExample;
259                digester.update_with_string(format!("enum {} (variants = {})", enum_name, #variant_count));
260                #serialized_variants
261                digester.create_child()
262            }
263        }
264    }).into()
265}
266
267#[cfg(RUSTC_WITH_SPECIALIZATION)]
268#[proc_macro_derive(AbiEnumVisitor)]
269pub fn derive_abi_enum_visitor(item: TokenStream) -> TokenStream {
270    let item = parse_macro_input!(item as Item);
271
272    match item {
273        Item::Enum(input) => do_derive_abi_enum_visitor(input),
274        _ => Error::new_spanned(item, "AbiEnumVisitor not applicable; only for enum")
275            .to_compile_error()
276            .into(),
277    }
278}
279
280#[cfg(RUSTC_WITH_SPECIALIZATION)]
281fn quote_for_test(
282    test_mod_ident: &Ident,
283    type_name: &Ident,
284    expected_digest: &str,
285) -> TokenStream2 {
286    // escape from nits.sh...
287    let p = Ident::new(&("ep".to_owned() + "rintln"), Span::call_site());
288    quote! {
289        #[cfg(test)]
290        mod #test_mod_ident {
291            use super::*;
292            use ::trezoa_frozen_abi::abi_example::{AbiExample, AbiEnumVisitor};
293
294            #[test]
295            fn test_abi_digest() {
296                ::trezoa_logger::setup();
297                let mut digester = ::trezoa_frozen_abi::abi_digester::AbiDigester::create();
298                let example = <#type_name>::example();
299                let result = <_>::visit_for_abi(&&example, &mut digester);
300                let mut hash = digester.finalize();
301                // pretty-print error
302                if result.is_err() {
303                    ::log::error!("digest error: {:#?}", result);
304                }
305                result.unwrap();
306                let actual_digest = format!("{}", hash);
307                if ::std::env::var("TREZOA_ABI_BULK_UPDATE").is_ok() {
308                    if #expected_digest != actual_digest {
309                        #p!("sed -i -e 's/{}/{}/g' $(git grep --files-with-matches frozen_abi)", #expected_digest, hash);
310                    }
311                    ::log::warn!("Not testing the abi digest under TREZOA_ABI_BULK_UPDATE!");
312                } else {
313                    if let Ok(dir) = ::std::env::var("TREZOA_ABI_DUMP_DIR") {
314                        assert_eq!(#expected_digest, actual_digest, "Possibly ABI changed? Examine the diff in TREZOA_ABI_DUMP_DIR!: \n$ diff -u {}/*{}* {}/*{}*", dir, #expected_digest, dir, actual_digest);
315                    } else {
316                        assert_eq!(#expected_digest, actual_digest, "Possibly ABI changed? Confirm the diff by rerunning before and after this test failed with TREZOA_ABI_DUMP_DIR!");
317                    }
318                }
319            }
320        }
321    }
322}
323
324#[cfg(RUSTC_WITH_SPECIALIZATION)]
325fn test_mod_name(type_name: &Ident) -> Ident {
326    Ident::new(&format!("{type_name}_frozen_abi"), Span::call_site())
327}
328
329#[cfg(RUSTC_WITH_SPECIALIZATION)]
330fn frozen_abi_type_alias(input: ItemType, expected_digest: &str) -> TokenStream {
331    let type_name = &input.ident;
332    let test = quote_for_test(&test_mod_name(type_name), type_name, expected_digest);
333    let result = quote! {
334        #input
335        #test
336    };
337    result.into()
338}
339
340#[cfg(RUSTC_WITH_SPECIALIZATION)]
341fn frozen_abi_struct_type(input: ItemStruct, expected_digest: &str) -> TokenStream {
342    let type_name = &input.ident;
343    let test = quote_for_test(&test_mod_name(type_name), type_name, expected_digest);
344    let result = quote! {
345        #input
346        #test
347    };
348    result.into()
349}
350
351#[cfg(RUSTC_WITH_SPECIALIZATION)]
352fn quote_sample_variant(
353    type_name: &Ident,
354    ty_generics: &syn::TypeGenerics,
355    variant: &Variant,
356) -> TokenStream2 {
357    let variant_name = &variant.ident;
358    let variant = &variant.fields;
359    if *variant == Fields::Unit {
360        quote! {
361            let sample_variant: #type_name #ty_generics = #type_name::#variant_name;
362        }
363    } else if let Fields::Unnamed(variant_fields) = variant {
364        let mut fields = quote! {};
365        for field in &variant_fields.unnamed {
366            if !(field.ident.is_none() && field.colon_token.is_none()) {
367                unimplemented!();
368            }
369            let ty = &field.ty;
370            fields.extend(quote! {
371                <#ty>::example(),
372            });
373        }
374        quote! {
375            let sample_variant: #type_name #ty_generics = #type_name::#variant_name(#fields);
376        }
377    } else if let Fields::Named(variant_fields) = variant {
378        let mut fields = quote! {};
379        for field in &variant_fields.named {
380            if field.ident.is_none() || field.colon_token.is_none() {
381                unimplemented!();
382            }
383            let field_type_name = &field.ty;
384            let field_name = &field.ident;
385            fields.extend(quote! {
386                #field_name: <#field_type_name>::example(),
387            });
388        }
389        quote! {
390            let sample_variant: #type_name #ty_generics = #type_name::#variant_name{#fields};
391        }
392    } else {
393        unimplemented!("variant: {:?}", variant)
394    }
395}
396
397#[cfg(RUSTC_WITH_SPECIALIZATION)]
398fn frozen_abi_enum_type(input: ItemEnum, expected_digest: &str) -> TokenStream {
399    let type_name = &input.ident;
400    let test = quote_for_test(&test_mod_name(type_name), type_name, expected_digest);
401    let result = quote! {
402        #input
403        #test
404    };
405    result.into()
406}
407
408#[cfg(RUSTC_WITH_SPECIALIZATION)]
409#[proc_macro_attribute]
410pub fn frozen_abi(attrs: TokenStream, item: TokenStream) -> TokenStream {
411    let mut expected_digest: Option<String> = None;
412    let attrs_parser = syn::meta::parser(|meta| {
413        if meta.path.is_ident("digest") {
414            expected_digest = Some(meta.value()?.parse::<LitStr>()?.value());
415            Ok(())
416        } else {
417            Err(meta.error("unsupported \"frozen_abi\" property"))
418        }
419    });
420    parse_macro_input!(attrs with attrs_parser);
421
422    let Some(expected_digest) = expected_digest else {
423        return Error::new_spanned(
424            TokenStream2::from(item),
425            "the required \"digest\" = ... attribute is missing.",
426        )
427        .to_compile_error()
428        .into();
429    };
430
431    let item = parse_macro_input!(item as Item);
432    match item {
433        Item::Struct(input) => frozen_abi_struct_type(input, &expected_digest),
434        Item::Enum(input) => frozen_abi_enum_type(input, &expected_digest),
435        Item::Type(input) => frozen_abi_type_alias(input, &expected_digest),
436        _ => Error::new_spanned(
437            item,
438            "frozen_abi isn't applicable; only for struct, enum and type",
439        )
440        .to_compile_error()
441        .into(),
442    }
443}