iter_all_proc/
lib.rs

1use proc_macro::TokenStream;
2use quote::{format_ident, quote, ToTokens};
3use syn::{parse_macro_input, Data, DeriveInput, Fields, Type};
4
5#[proc_macro_derive(IterAll)]
6pub fn iter_all_macro(input: TokenStream) -> TokenStream {
7    let input = parse_macro_input!(input as DeriveInput);
8
9    let enum_name = &input.ident;
10
11    let Data::Enum(ref data_enum) = input.data else {
12        panic!("IterAll can only be derived for enums");
13    };
14
15    let variants = data_enum.variants.iter().map(|variant| {
16        let variant_name = &variant.ident;
17
18        if let Fields::Unit = variant.fields {
19            return quote! {
20                action(#enum_name::#variant_name);
21            };
22        }
23
24        quote! {
25            action(#enum_name::#variant_name(Default::default()));
26        }
27    });
28
29    let variants = quote! {
30        #(#variants)*
31    };
32
33    let variants3 = data_enum.variants.iter().map(|variant| {
34        let variant_name = &variant.ident;
35
36        let uppercase_ident = format_ident!("{}", variant_name.to_string().to_uppercase());
37
38        if let Fields::Unit = variant.fields {
39            return quote! {
40                pub const #uppercase_ident: #enum_name = #enum_name::#variant;
41            };
42        }
43
44        let Fields::Unnamed(ref unnamed) = variant.fields else {
45            panic!();
46        };
47
48        let ty: &Type = &unnamed.unnamed.first().unwrap().ty;
49
50        let ty_str = ty.to_owned().to_token_stream().to_string();
51
52        match (ty_str.find('<'), ty_str.rfind('>')) {
53            (Some(start), Some(end)) => {
54                let generics = &ty_str[start + 1..end].trim();
55
56                let only_type = &ty_str[0..start - 1].trim();
57
58                let only_type = format_ident!("{}", only_type);
59
60                use std::str::FromStr;
61
62                let generics: proc_macro2::TokenStream =
63                    TokenStream::from_str(generics).unwrap().to_owned().into();
64
65                quote! {
66                    pub const #uppercase_ident: #ty = #only_type::<#generics>::new();
67                }
68            }
69            _ => quote! {
70                pub const #uppercase_ident: #ty = #ty::new();
71            },
72        }
73    });
74
75    let variants3 = quote! {
76        #(#variants3)*
77    };
78
79    quote! {
80        impl iter_all::IterAll for #enum_name {
81            fn iter_all(mut action: impl FnMut(Self)) {
82                #variants
83            }
84        }
85
86        impl #enum_name {
87            #variants3
88        }
89    }
90    .into()
91}
iter_all_proc/lib.rs

iter_all_proc/
lib.rs
