1
use proc_macro::TokenStream;
2
use quote::quote;
3
use syn::{FnArg, GenericArgument, ItemFn, PathArguments, Type, parse_macro_input};
4

            
5
/// Creates a named rule wrapper function for tree-morph rules.
6
///
7
/// This macro transforms a rule function into a `NamedRule`.
8
/// by wrapping it in a `NamedRule`.
9
///
10
///  ```rust
11
/// use tree_morph::prelude::*;
12
/// use tree_morph_macros::named_rule;
13
/// use uniplate::Uniplate;
14
///
15
/// #[derive(Debug, Clone, PartialEq, Eq, Uniplate)]
16
/// #[uniplate()]
17
/// enum Expr {
18
///     # Add(Box<Expr>, Box<Expr>),
19
///    // Snip
20
/// }
21
///
22
/// struct Meta;
23
///
24
/// #[named_rule("CustomName")]
25
/// fn my_rule(_: &mut Commands<Expr, Meta>, expr: &Expr, _: &Meta) -> Option<Expr> {
26
///     /// rule implementation
27
///     # None
28
/// }
29
/// ```
30
///
31
/// The above function will be transformed into a `NamedRule` and
32
/// have its name set to "CustomName". If not specified, the name will simply be
33
/// the function identifier. In the above case it will be "my_rule".
34
///
35
/// The original function logic is preserved in a private helper function.
36
#[proc_macro_attribute]
37
7
pub fn named_rule(attr: TokenStream, item: TokenStream) -> TokenStream {
38
7
    let function = parse_macro_input!(item as ItemFn);
39
7
    let name = &function.sig.ident;
40
7
    let vis = &function.vis;
41

            
42
7
    let rule_name = if attr.is_empty() {
43
        name.to_string()
44
    } else {
45
7
        let name_lit: syn::LitStr = parse_macro_input!(attr as syn::LitStr);
46
7
        name_lit.value()
47
    };
48

            
49
    // Create a private helper function with the original implementation
50
7
    let helper_name = syn::Ident::new(&format!("__{name}_impl"), name.span());
51
7
    let mut helper_function = function.clone();
52
7
    helper_function.sig.ident = helper_name.clone();
53
7
    helper_function.vis = syn::Visibility::Inherited;
54

            
55
7
    let commmand_param = &function.sig.inputs[0];
56
7
    let (type_t, type_m) = extract_commands_types(commmand_param);
57

            
58
7
    let expanded = quote! {
59
        #helper_function
60

            
61
        #[allow(non_upper_case_globals)]
62
        #vis const #name: ::tree_morph::prelude::NamedRule<::tree_morph::prelude::RuleFn<#type_t, #type_m>> =
63
          ::tree_morph::prelude::NamedRule::new(#rule_name, #helper_name as ::tree_morph::prelude::RuleFn<#type_t, #type_m>);
64
    };
65

            
66
7
    TokenStream::from(expanded)
67
7
}
68

            
69
7
fn extract_commands_types(param: &syn::FnArg) -> (syn::Type, syn::Type) {
70
7
    let FnArg::Typed(pat_type) = param else {
71
        panic!("Expected typed parameter");
72
    };
73

            
74
7
    let Type::Reference(type_ref) = &*pat_type.ty else {
75
        panic!("Expected reference type");
76
    };
77

            
78
7
    let Type::Path(type_path) = &*type_ref.elem else {
79
        panic!("Expected path type");
80
    };
81

            
82
7
    let segment = type_path.path.segments.last().unwrap();
83

            
84
7
    let PathArguments::AngleBracketed(args) = &segment.arguments else {
85
        panic!("Commands must have type parameters");
86
    };
87

            
88
7
    let GenericArgument::Type(t) = &args.args[0] else {
89
        panic!("First argument must be a type");
90
    };
91

            
92
7
    let GenericArgument::Type(m) = &args.args[1] else {
93
        panic!("Second argument must be a type");
94
    };
95

            
96
7
    (t.clone(), m.clone())
97
7
}