1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
extern crate proc_macro;
use proc_macro::TokenStream;
use quote::quote;
//use syn::punctuated::Punctuated;
use syn::{parse_macro_input, AttributeArgs, ItemFn, NestedMeta};

#[proc_macro_attribute]
pub fn modifiers(args: TokenStream, input: TokenStream) -> TokenStream {
    let args = parse_macro_input!(args as AttributeArgs);
    let func = parse_macro_input!(input as ItemFn);

    // Transform attribute arguments to strings
    let mut modifiers: Vec<(String, Vec<String>)> = Vec::new();
    for arg in args {
        if let NestedMeta::Lit(syn::Lit::Str(lit)) = arg {
            let val = lit.value();
            let parts: Vec<_> = val.split("@").collect();
            let func_name = parts[0].to_string();
            let params = if parts.len() > 1 {
                parts[1..]
                    .join("@")
                    .split(',')
                    .map(|s| s.trim().to_string())
                    .collect()
            } else {
                Vec::new()
            };
            modifiers.push((func_name, params));
        }
    }

    // The function name and parameters
    let fn_name = &func.sig.ident;
    let fn_params = &func.sig.inputs;

    // Extract statements from the original block
    let fn_stmts = &func.block.stmts;

    // Generate each modifier check and function call
    let modifier_checks: Vec<proc_macro2::TokenStream> = modifiers
        .iter()
        .map(|(modi, params)| {
            let modi_ident: syn::Ident = syn::Ident::new(&modi, proc_macro2::Span::call_site());
            let params_tokens: Vec<proc_macro2::TokenStream> = params
                .iter()
                .map(|p| {
                    let param_token: proc_macro2::TokenStream = p.parse().unwrap();
                    param_token
                })
                .collect();
            quote! {
                let r: Result<(), String> = #modi_ident(#(#params_tokens),*);
                if let Err(e) = r {
                    ic_cdk::api::call::reject(&e);
                    return;
                }
            }
        })
        .collect();

    let expanded = quote! {
        fn #fn_name(#fn_params) {
            #(#modifier_checks)*
            #(#fn_stmts)*
        }
    };

    TokenStream::from(expanded)
}