1
use proc_macro::TokenStream;
2
use quote::quote;
3
use syn::{ItemFn, Token, parse_macro_input};
4

            
5
enum TriggerType {
6
    EntityType(syn::Ident),
7
    Function(syn::Ident),
8
}
9

            
10
struct ScriptArgs {
11
    trigger: TriggerType,
12
    output_size: Option<u32>,
13
}
14

            
15
impl syn::parse::Parse for ScriptArgs {
16
    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
17
        let mut trigger = None;
18
        let mut output_size = None;
19

            
20
        while !input.is_empty() {
21
            let ident: syn::Ident = input.parse()?;
22
            input.parse::<Token![=]>()?;
23

            
24
            match ident.to_string().as_str() {
25
                "trigger" => {
26
                    let value: syn::Ident = input.parse()?;
27
                    trigger = Some(TriggerType::EntityType(value));
28
                }
29
                "trigger_fn" => {
30
                    let value: syn::Ident = input.parse()?;
31
                    trigger = Some(TriggerType::Function(value));
32
                }
33
                "output_size" => {
34
                    let lit: syn::LitInt = input.parse()?;
35
                    output_size = Some(lit.base10_parse()?);
36
                }
37
                _ => {
38
                    return Err(syn::Error::new(ident.span(), "unknown attribute"));
39
                }
40
            }
41

            
42
            if input.peek(Token![,]) {
43
                input.parse::<Token![,]>()?;
44
            }
45
        }
46

            
47
        let trigger = trigger.ok_or_else(|| {
48
            syn::Error::new(
49
                proc_macro2::Span::call_site(),
50
                "missing `trigger` or `trigger_fn` argument",
51
            )
52
        })?;
53

            
54
        Ok(ScriptArgs {
55
            trigger,
56
            output_size,
57
        })
58
    }
59
}
60

            
61
/// Attribute macro for defining Nomisync WASM scripts.
62
///
63
/// # Arguments
64
/// - `trigger`: The entity type that triggers this script (Transaction, Split, Tag, etc.)
65
/// - `trigger_fn`: A function that receives `&Context` and returns `bool`
66
/// - `output_size` (optional): Size of output buffer in bytes (default: 4096)
67
///
68
/// # Example with entity type trigger
69
/// ```ignore
70
/// use scripting_sdk::prelude::*;
71
///
72
/// #[script(trigger = Transaction)]
73
/// fn add_tag(ctx: &mut Context) -> Result<()> {
74
///     ctx.tag_primary("processed", "true")?;
75
///     Ok(())
76
/// }
77
/// ```
78
///
79
/// # Example with function trigger
80
/// ```ignore
81
/// use scripting_sdk::prelude::*;
82
///
83
/// fn is_groceries(ctx: &Context) -> bool {
84
///     // Check if transaction has tag "note" = "groceries"
85
///     for tag in ctx.tags_for(ctx.primary_entity_idx()).flatten() {
86
///         if tag.name().ok() == Some("note") && tag.value().ok() == Some("groceries") {
87
///             return true;
88
///         }
89
///     }
90
///     false
91
/// }
92
///
93
/// #[script(trigger_fn = is_groceries)]
94
/// fn categorize(ctx: &mut Context) -> Result<()> {
95
///     // Tag all splits with category
96
///     for split in ctx.splits_for(ctx.primary_entity_idx()).flatten() {
97
///         ctx.create_tag(split.parent_idx(), "category", "groceries")?;
98
///     }
99
///     Ok(())
100
/// }
101
/// ```
102
#[proc_macro_attribute]
103
pub fn script(attr: TokenStream, item: TokenStream) -> TokenStream {
104
    let args = parse_macro_input!(attr as ScriptArgs);
105
    let input_fn = parse_macro_input!(item as ItemFn);
106

            
107
    let fn_name = &input_fn.sig.ident;
108
    let fn_block = &input_fn.block;
109

            
110
    let context_creation = if let Some(size) = args.output_size {
111
        quote! { Context::with_output_size(#size)? }
112
    } else {
113
        quote! { Context::new()? }
114
    };
115

            
116
    let should_apply_body = match &args.trigger {
117
        TriggerType::EntityType(entity_type) => quote! {
118
            let Ok(ctx) = Context::new() else {
119
                return 0;
120
            };
121
            let Ok(entity_type) = ctx.primary_entity_type() else {
122
                return 0;
123
            };
124
            if entity_type == EntityType::#entity_type {
125
                1
126
            } else {
127
                0
128
            }
129
        },
130
        TriggerType::Function(trigger_fn) => quote! {
131
            let Ok(ctx) = Context::new() else {
132
                return 0;
133
            };
134
            if #trigger_fn(&ctx) {
135
                1
136
            } else {
137
                0
138
            }
139
        },
140
    };
141

            
142
    let expanded = quote! {
143
        use core::panic::PanicInfo;
144

            
145
        #[panic_handler]
146
        fn panic(_info: &PanicInfo) -> ! {
147
            loop {}
148
        }
149

            
150
        #[unsafe(no_mangle)]
151
        pub unsafe extern "C" fn should_apply() -> i32 {
152
            #should_apply_body
153
        }
154

            
155
        #[unsafe(no_mangle)]
156
        pub unsafe extern "C" fn process() {
157
            let _ = __script_impl();
158
        }
159

            
160
        fn __script_impl() -> Result<()> {
161
            let mut ctx = #context_creation;
162
            let result = #fn_name(&mut ctx);
163
            ctx.finalize()?;
164
            result
165
        }
166

            
167
        fn #fn_name(ctx: &mut Context) -> Result<()>
168
            #fn_block
169
    };
170

            
171
    TokenStream::from(expanded)
172
}