Skip to main content

async_nodes_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::{format_ident, quote};
3use syn::{FnArg, ItemFn, Pat, parse_macro_input};
4
5/// Formats standard snake_case function names to PascalCase for the generated Struct.
6fn to_pascal_case(s: &str) -> String {
7    s.split('_')
8        .filter(|w| !w.is_empty())
9        .map(|w| {
10            let mut c = w.chars();
11            c.next().map_or_else(String::new, |first| {
12                first.to_uppercase().collect::<String>() + c.as_str()
13            })
14        })
15        .collect()
16}
17
18fn is_option_type(ty: &syn::Type) -> bool {
19    if let syn::Type::Path(type_path) = ty {
20        if let Some(segment) = type_path.path.segments.last() {
21            return segment.ident == "Option";
22        }
23    }
24    false
25}
26
27/// Helper to detect, extract, and optionally parse a key type from `#[context(KeyType)]`.
28fn has_and_remove_context_attr(pat_type: &mut syn::PatType) -> Option<Option<syn::Type>> {
29    let mut found = None;
30    pat_type.attrs.retain(|attr| {
31        if attr.path().is_ident("context") {
32            // Check if there is an argument like #[context(HeavyResource)]
33            if let syn::Meta::List(meta_list) = &attr.meta {
34                if let Ok(key_ty) = meta_list.parse_args::<syn::Type>() {
35                    found = Some(Some(key_ty));
36                } else {
37                    found = Some(None);
38                }
39            } else {
40                found = Some(None);
41            }
42            false // Discard attribute
43        } else {
44            true // Keep other attributes
45        }
46    });
47    found
48}
49
50/// Unpacks the dependency as a cloned, owned type.
51/// Rejects reference parameters to prevent lock-guards from surviving across async suspension points.
52fn parse_context_arg(
53    arg_name: &syn::Ident,
54    arg_ty: &syn::Type,
55    key_override: Option<syn::Type>,
56) -> Result<proc_macro2::TokenStream, syn::Error> {
57    if let syn::Type::Reference(type_ref) = arg_ty {
58        return Err(syn::Error::new_spanned(
59            type_ref,
60            "Context parameters must be owned types (e.g., implementing `Clone` or wrapped in an `Arc`). Borrowed references are not permitted.",
61        ));
62    }
63
64    let key_ty = key_override.as_ref().unwrap_or(arg_ty);
65    Ok(quote! {
66        let #arg_name = {
67            let guard = run.ctx.get::<#key_ty>().ok_or_else(|| {
68                ::async_nodes::error::BspError::ValidationError(format!(
69                    "Dependency injection failed: TypeMapKey `{}` not found in BspContext",
70                    ::std::any::type_name::<#key_ty>()
71                ))
72            })?;
73            (*guard).clone()
74        };
75    })
76}
77
78fn get_inner_type_if_result(ty: &syn::Type) -> &syn::Type {
79    if let syn::Type::Path(type_path) = ty {
80        if let Some(segment) = type_path.path.segments.last() {
81            if segment.ident == "Result" {
82                if let syn::PathArguments::AngleBracketed(args) = &segment.arguments {
83                    if let Some(syn::GenericArgument::Type(inner_ty)) = args.args.first() {
84                        return inner_ty;
85                    }
86                }
87            }
88        }
89    }
90    ty
91}
92
93/// Procedural attribute transforming an asynchronous function into an explicit `BspNode`.
94#[proc_macro_attribute]
95pub fn node(attr: TokenStream, item: TokenStream) -> TokenStream {
96    let node_name_override = if !attr.is_empty() {
97        let lit = parse_macro_input!(attr as syn::LitStr);
98        Some(lit.value())
99    } else {
100        None
101    };
102
103    let mut input_fn = parse_macro_input!(item as ItemFn);
104    let fn_name = &input_fn.sig.ident;
105    let fn_vis = &input_fn.vis;
106    let fn_attrs = &input_fn.attrs;
107
108    let struct_name = format_ident!("{}Node", to_pascal_case(&fn_name.to_string()));
109    let ext_trait_name = format_ident!("{struct_name}Ports");
110    let node_name_str = node_name_override.unwrap_or_else(|| fn_name.to_string());
111
112    let mut input_ports = Vec::new();
113    let mut unpack_inputs = Vec::new();
114    let mut context_unpacks = Vec::new();
115    let mut call_args = Vec::new();
116    let mut port_idx: usize = 0; // Explicitly compile-time usize
117
118    let mut trait_signatures = Vec::new();
119    let mut handle_methods = Vec::new();
120
121    for arg in input_fn.sig.inputs.iter_mut() {
122        if let FnArg::Typed(pat_type) = arg {
123            if let Pat::Ident(pat_ident) = &*pat_type.pat {
124                let arg_name = pat_ident.ident.clone();
125                let arg_ty = (*pat_type.ty).clone();
126                call_args.push(quote! { #arg_name });
127
128                if let Some(key_override) = has_and_remove_context_attr(pat_type) {
129                    match parse_context_arg(&arg_name, &arg_ty, key_override) {
130                        Ok(tokens) => context_unpacks.push(tokens),
131                        Err(err) => return err.into_compile_error().into(),
132                    }
133                } else {
134                    let idx = port_idx;
135                    port_idx += 1;
136
137                    trait_signatures.push(quote! {
138                        #[must_use]
139                        fn #arg_name(&self) -> ::async_nodes::graph::InputPin;
140                    });
141
142                    handle_methods.push(quote! {
143                        fn #arg_name(&self) -> ::async_nodes::graph::InputPin {
144                            ::async_nodes::graph::InputPin {
145                                node_id: self.id,
146                                port: #idx,
147                            }
148                        }
149                    });
150
151                    input_ports.push(quote! { #idx as usize });
152
153                    if is_option_type(&arg_ty) {
154                        unpack_inputs.push(quote! {
155                            let #arg_name = {
156                                if let Some(val) = run.take_input(#idx as usize) {
157                                    <#arg_ty as ::async_nodes::value::BspFromValue>::from_value(val, run.node_id, #idx as usize)?
158                                } else {
159                                    None
160                                }
161                            };
162                        });
163                    } else {
164                        unpack_inputs.push(quote! {
165                            let #arg_name = {
166                                let val = run.take_input(#idx as usize)
167                                    .ok_or_else(|| ::async_nodes::error::BspError::UnconnectedPort {
168                                        node_id: run.node_id,
169                                        port_index: #idx as usize,
170                                    })?;
171                                <#arg_ty as ::async_nodes::value::BspFromValue>::from_value(val, run.node_id, #idx as usize)?
172                            };
173                        });
174                    }
175                }
176            }
177        }
178    }
179
180    // Capture the return type path to delegate port count statically
181    let return_type = match &input_fn.sig.output {
182        syn::ReturnType::Default => quote! { () },
183        syn::ReturnType::Type(_, ty) => quote! { #ty },
184    };
185
186    match &input_fn.sig.output {
187        syn::ReturnType::Default => {}
188        syn::ReturnType::Type(_, ty) => {
189            let inner_ty = get_inner_type_if_result(ty);
190            if let syn::Type::Tuple(type_tuple) = inner_ty {
191                for idx in 0..type_tuple.elems.len() {
192                    let method_name = format_ident!("output_{}", idx);
193                    trait_signatures.push(quote! {
194                        #[must_use]
195                        fn #method_name(&self) -> ::async_nodes::graph::OutputPin;
196                    });
197                    handle_methods.push(quote! {
198                        fn #method_name(&self) -> ::async_nodes::graph::OutputPin {
199                            ::async_nodes::graph::OutputPin {
200                                node_id: self.id,
201                                port: #idx,
202                            }
203                        }
204                    });
205                }
206            } else if let syn::Type::Path(path_type) = inner_ty {
207                if !path_type.path.is_ident("()") {
208                    trait_signatures.push(quote! {
209                        #[must_use]
210                        fn output(&self) -> ::async_nodes::graph::OutputPin;
211                    });
212                    handle_methods.push(quote! {
213                        fn output(&self) -> ::async_nodes::graph::OutputPin {
214                            ::async_nodes::graph::OutputPin {
215                                node_id: self.id,
216                                port: 0,
217                            }
218                        }
219                    });
220                }
221            } else {
222                trait_signatures.push(quote! {
223                    #[must_use]
224                    fn output(&self) -> ::async_nodes::graph::OutputPin;
225                });
226                handle_methods.push(quote! {
227                    fn output(&self) -> ::async_nodes::graph::OutputPin {
228                        ::async_nodes::graph::OutputPin {
229                            node_id: self.id,
230                            port: 0,
231                        }
232                    }
233                });
234            }
235        }
236    }
237
238    let expanded = quote! {
239        #(#fn_attrs)*
240        #[allow(clippy::unused_async, clippy::needless_pass_by_value)]
241        #fn_vis #input_fn
242
243        #[derive(Clone, Default)]
244        #fn_vis struct #struct_name;
245
246        impl #struct_name {
247            pub fn new() -> Self { Self }
248        }
249
250        #fn_vis trait #ext_trait_name {
251            #(#trait_signatures)*
252        }
253
254        impl #ext_trait_name for ::async_nodes::graph::TypedNodeHandle<#struct_name> {
255            #(#handle_methods)*
256        }
257
258        #[::async_nodes::async_trait::async_trait]
259        impl<S: Send> ::async_nodes::node::BspNode<S> for #struct_name {
260            fn name(&self) -> &'static str { #node_name_str }
261            fn input_ports(&self) -> ::std::vec::Vec<::async_nodes::graph::PortIndex> { ::std::vec![#(#input_ports),*] }
262            fn output_ports(&self) -> ::std::vec::Vec<::async_nodes::graph::PortIndex> { <#return_type as ::async_nodes::value::BspOutputPacker>::output_ports() }
263            fn is_dynamic_output(&self) -> bool { <#return_type as ::async_nodes::value::BspOutputPacker>::is_dynamic() }
264            fn activation(&self) -> ::async_nodes::node::Activation { ::async_nodes::node::Activation::All }
265
266            async fn execute(
267                &self,
268                run: &mut ::async_nodes::run::RunContext<'_, S>,
269            ) -> ::std::result::Result<::async_nodes::node::NodeOutcome, ::async_nodes::error::BspError> {
270                // 1. Unpack all input ports safely using static trait resolution
271                #(#context_unpacks)*
272                #(#unpack_inputs)*
273
274                // 2. Execute the user's native async logic
275                let result = #fn_name(#(#call_args),*).await;
276
277                // 3. Serialize outputs or map execution faults
278                <#return_type as ::async_nodes::value::BspOutputPacker>::pack(result, run.node_id, #node_name_str)
279            }
280        }
281    };
282
283    TokenStream::from(expanded)
284}
285
286/// Procedural attribute transforming a synchronous, blocking function into a `BspNode`.
287#[proc_macro_attribute]
288pub fn node_blocking(attr: TokenStream, item: TokenStream) -> TokenStream {
289    let node_name_override = if !attr.is_empty() {
290        let lit = parse_macro_input!(attr as syn::LitStr);
291        Some(lit.value())
292    } else {
293        None
294    };
295
296    let mut input_fn = parse_macro_input!(item as ItemFn);
297    let fn_name = &input_fn.sig.ident;
298    let fn_vis = &input_fn.vis;
299    let fn_attrs = &input_fn.attrs;
300
301    let struct_name = format_ident!("{}Node", to_pascal_case(&fn_name.to_string()));
302    let ext_trait_name = format_ident!("{struct_name}Ports");
303    let node_name_str = node_name_override.unwrap_or_else(|| fn_name.to_string());
304
305    let mut input_ports = Vec::new();
306    let mut unpack_inputs = Vec::new();
307    let mut context_unpacks = Vec::new();
308    let mut call_args = Vec::new();
309    let mut port_idx: usize = 0; // Explicitly compile-time usize
310
311    let mut trait_signatures = Vec::new();
312    let mut handle_methods = Vec::new();
313
314    for arg in input_fn.sig.inputs.iter_mut() {
315        if let FnArg::Typed(pat_type) = arg {
316            if let Pat::Ident(pat_ident) = &*pat_type.pat {
317                let arg_name = pat_ident.ident.clone();
318                let arg_ty = (*pat_type.ty).clone();
319                call_args.push(quote! { #arg_name });
320
321                if let Some(key_override) = has_and_remove_context_attr(pat_type) {
322                    match parse_context_arg(&arg_name, &arg_ty, key_override) {
323                        Ok(tokens) => context_unpacks.push(tokens),
324                        Err(err) => return err.into_compile_error().into(),
325                    }
326                } else {
327                    let idx = port_idx;
328                    port_idx += 1;
329
330                    trait_signatures.push(quote! {
331                        #[must_use]
332                        fn #arg_name(&self) -> ::async_nodes::graph::InputPin;
333                    });
334
335                    handle_methods.push(quote! {
336                        fn #arg_name(&self) -> ::async_nodes::graph::InputPin {
337                            ::async_nodes::graph::InputPin {
338                                node_id: self.id,
339                                port: #idx,
340                            }
341                        }
342                    });
343
344                    input_ports.push(quote! { #idx as usize });
345
346                    if is_option_type(&arg_ty) {
347                        unpack_inputs.push(quote! {
348                            let #arg_name = {
349                                if let Some(val) = run.take_input(#idx as usize) {
350                                    <#arg_ty as ::async_nodes::value::BspFromValue>::from_value(val, run.node_id, #idx as usize)?
351                                } else {
352                                    None
353                                }
354                            };
355                        });
356                    } else {
357                        unpack_inputs.push(quote! {
358                            let #arg_name = {
359                                let val = run.take_input(#idx as usize)
360                                    .ok_or_else(|| ::async_nodes::error::BspError::UnconnectedPort {
361                                        node_id: run.node_id,
362                                        port_index: #idx as usize,
363                                    })?;
364                                <#arg_ty as ::async_nodes::value::BspFromValue>::from_value(val, run.node_id, #idx as usize)?
365                            };
366                        });
367                    }
368                }
369            }
370        }
371    }
372
373    let return_type = match &input_fn.sig.output {
374        syn::ReturnType::Default => quote! { () },
375        syn::ReturnType::Type(_, ty) => quote! { #ty },
376    };
377
378    match &input_fn.sig.output {
379        syn::ReturnType::Default => {}
380        syn::ReturnType::Type(_, ty) => {
381            let inner_ty = get_inner_type_if_result(ty);
382            if let syn::Type::Tuple(type_tuple) = inner_ty {
383                for idx in 0..type_tuple.elems.len() {
384                    let method_name = format_ident!("output_{}", idx);
385                    trait_signatures.push(quote! {
386                        #[must_use]
387                        fn #method_name(&self) -> ::async_nodes::graph::OutputPin;
388                    });
389                    handle_methods.push(quote! {
390                        fn #method_name(&self) -> ::async_nodes::graph::OutputPin {
391                            ::async_nodes::graph::OutputPin {
392                                node_id: self.id,
393                                port: #idx,
394                            }
395                        }
396                    });
397                }
398            } else if let syn::Type::Path(path_type) = inner_ty {
399                if !path_type.path.is_ident("()") {
400                    trait_signatures.push(quote! {
401                        #[must_use]
402                        fn output(&self) -> ::async_nodes::graph::OutputPin;
403                    });
404                    handle_methods.push(quote! {
405                        fn output(&self) -> ::async_nodes::graph::OutputPin {
406                            ::async_nodes::graph::OutputPin {
407                                node_id: self.id,
408                                port: 0,
409                            }
410                        }
411                    });
412                }
413            } else {
414                trait_signatures.push(quote! {
415                    #[must_use]
416                    fn output(&self) -> ::async_nodes::graph::OutputPin;
417                });
418                handle_methods.push(quote! {
419                    fn output(&self) -> ::async_nodes::graph::OutputPin {
420                        ::async_nodes::graph::OutputPin {
421                            node_id: self.id,
422                            port: 0,
423                        }
424                    }
425                });
426            }
427        }
428    }
429
430    let expanded = quote! {
431        #(#fn_attrs)*
432        #[allow(clippy::needless_pass_by_value)]
433        #fn_vis #input_fn
434
435        #[derive(Clone, Default)]
436        #fn_vis struct #struct_name;
437
438        impl #struct_name {
439            pub fn new() -> Self { Self }
440        }
441
442        #fn_vis trait #ext_trait_name {
443            #(#trait_signatures)*
444        }
445
446        impl #ext_trait_name for ::async_nodes::graph::TypedNodeHandle<#struct_name> {
447            #(#handle_methods)*
448        }
449
450        #[::async_nodes::async_trait::async_trait]
451        impl<S: Send> ::async_nodes::node::BspNode<S> for #struct_name {
452            fn name(&self) -> &'static str { #node_name_str }
453            fn input_ports(&self) -> ::std::vec::Vec<::async_nodes::graph::PortIndex> { ::std::vec![#(#input_ports),*] }
454            fn output_ports(&self) -> ::std::vec::Vec<::async_nodes::graph::PortIndex> { <#return_type as ::async_nodes::value::BspOutputPacker>::output_ports() }
455            fn is_dynamic_output(&self) -> bool { <#return_type as ::async_nodes::value::BspOutputPacker>::is_dynamic() }
456            fn activation(&self) -> ::async_nodes::node::Activation { ::async_nodes::node::Activation::All }
457
458            async fn execute(
459                &self,
460                run: &mut ::async_nodes::run::RunContext<'_, S>,
461            ) -> ::std::result::Result<::async_nodes::node::NodeOutcome, ::async_nodes::error::BspError> {
462                // 1. Unpack all input ports safely using static trait resolution
463                #(#context_unpacks)*
464                #(#unpack_inputs)*
465
466                // 2. Offload user's synchronous logic to spawn_blocking
467                let result = ::async_nodes::tokio::task::spawn_blocking(move || {
468                    #fn_name(#(#call_args),*)
469                })
470                .await
471                .map_err(|err| ::async_nodes::error::BspError::NodeExecutionFailed {
472                    node_id: run.node_id,
473                    node_name: #node_name_str.to_string(),
474                    message: ::std::string::ToString::to_string(&err),
475                })?;
476
477                // 3. Serialize outputs or map execution faults
478                <#return_type as ::async_nodes::value::BspOutputPacker>::pack(result, run.node_id, #node_name_str)
479            }
480        }
481    };
482
483    TokenStream::from(expanded)
484}
485
486/// Custom derive macro providing serialization support to arbitrary developer types.
487#[proc_macro_derive(BspType)]
488pub fn derive_bsp_type(input: TokenStream) -> TokenStream {
489    let input = parse_macro_input!(input as syn::DeriveInput);
490    let name = &input.ident;
491    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
492
493    let expanded = quote! {
494        impl #impl_generics ::async_nodes::value::BspFromValue for #name #ty_generics #where_clause {
495            fn from_value(
496                value: ::async_nodes::value::BspValue,
497                node_id: ::async_nodes::graph::NodeId,
498                port: ::async_nodes::graph::PortIndex,
499            ) -> ::std::result::Result<Self, ::async_nodes::error::BspError> {
500                let found_type = value.type_name();
501                let arc_t = value
502                    .downcast::<Self>()
503                    .map_err(|_| ::async_nodes::error::BspError::TypeMismatch {
504                        node_id,
505                        port_index: port,
506                        expected: ::std::any::type_name::<Self>(),
507                        found: found_type,
508                    })?;
509                Ok((*arc_t).clone())
510            }
511        }
512
513        impl #impl_generics ::async_nodes::value::BspToValue for #name #ty_generics #where_clause {
514            fn to_value(self) -> ::async_nodes::value::BspValue {
515                ::async_nodes::value::BspValue::new(self)
516            }
517        }
518
519        impl #impl_generics ::async_nodes::value::BspOutputPacker for #name #ty_generics #where_clause {
520            fn output_ports() -> ::std::vec::Vec<::async_nodes::graph::PortIndex> {
521                ::std::vec![0]
522            }
523            fn pack(
524                self,
525                _node_id: ::async_nodes::graph::NodeId,
526                _node_name: &str,
527            ) -> ::std::result::Result<::async_nodes::node::NodeOutcome, ::async_nodes::error::BspError> {
528                let mut outputs = ::indexmap::IndexMap::new();
529                outputs.insert(0, ::async_nodes::value::BspToValue::to_value(self));
530                ::std::result::Result::Ok(::async_nodes::node::NodeOutcome::Success(outputs))
531            }
532        }
533    };
534
535    TokenStream::from(expanded)
536}