shiv_macro_impl/
label.rs

1use proc_macro2::{Ident, Span, TokenStream};
2use quote::quote;
3use syn::{Data, DeriveInput, Fields, Path};
4
5pub fn derive_label(input: DeriveInput, shiv: Path, label: &str) -> proc_macro2::TokenStream {
6    let ident = Ident::new(label, Span::call_site());
7    let id = Ident::new(&format!("{label}Id"), Span::call_site());
8
9    let name = &input.ident;
10
11    let label_impl = label_fn_impl(&input, &id, &shiv);
12
13    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
14
15    quote! {
16        #[automatically_derived]
17        impl #impl_generics #shiv::schedule::#ident for #name #ty_generics #where_clause {
18            #[inline]
19            fn label(self) -> #shiv::schedule::#id {
20                #label_impl
21            }
22        }
23    }
24}
25
26fn label_fn_impl(input: &DeriveInput, id: &Ident, shiv: &Path) -> TokenStream {
27    match input.data {
28        Data::Struct(ref data) => {
29            if !matches!(data.fields, Fields::Unit) {
30                unimplemented!("Only unit structs are supported");
31            }
32
33            let name = input.ident.to_string();
34
35            quote! {
36                #shiv::schedule::#id::from_raw_parts::<Self>(#name, 0u32)
37            }
38        }
39        Data::Enum(ref data) => {
40            let name = input.ident.to_string();
41
42            let variants = data.variants.iter().enumerate().map(|(i, variant)| {
43                if !matches!(variant.fields, Fields::Unit) {
44                    unimplemented!("Only unit variants are supported");
45                }
46
47                let variant_ident = &variant.ident;
48                let variant_name = variant.ident.to_string();
49                let name = format!("{}::{}", name, variant_name);
50
51                let i = i as u32;
52
53                quote! {
54                    Self::#variant_ident => #shiv::schedule::#id::from_raw_parts::<Self>(#name, #i)
55                }
56            });
57
58            quote! {
59                match self {
60                    #(#variants),*
61                }
62            }
63        }
64        _ => unimplemented!("Only structs and enums are supported"),
65    }
66}