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}