directed_stage_macro/
lib.rs

1mod parse;
2
3use parse::*;
4use proc_macro::TokenStream;
5use proc_macro_error::proc_macro_error;
6use proc_macro2::Span;
7use proc_macro2_diagnostics::Diagnostic;
8use quote::quote_spanned;
9use syn::{ItemFn, Token, parse_macro_input};
10
11/// A macro that wraps a function with the standardized interface:
12/// TODO: More in-depth docs here
13#[proc_macro_attribute]
14#[proc_macro_error]
15pub fn stage(attr: TokenStream, item: TokenStream) -> proc_macro::TokenStream {
16    let input_fn = parse_macro_input!(item as ItemFn);
17    let meta_args = parse_macro_input!(attr as StageArgs);
18    let stage_config = StageConfig::from_args(&input_fn, &meta_args).unwrap();
19    match generate_stage_impl(stage_config) {
20        Ok(tokens) => tokens.into(),
21        Err(diag) => diag.emit_as_expr_tokens().into(),
22    }
23}
24
25/// Get a type, squash the &
26fn true_type(ty: &syn::Type) -> &syn::Type {
27    if let syn::Type::Reference(ty) = ty {
28        &*ty.elem
29    } else {
30        ty
31    }
32}
33
34/// This code is used by the wrapper function - it downcasts type-erased
35/// function parameters so that the user-facing function can be called with
36/// concrete types.
37fn generate_extraction_code(
38    inputs: &[InputParam],
39    cache_strategy: (CacheStrategy, Span),
40) -> Result<Vec<proc_macro2::TokenStream>, Diagnostic> {
41    Ok(inputs.iter().map(|input| {
42        let arg_name = &input.clean_name;
43        let arg_type = true_type(&input.ty);
44        let clean_arg_name = input.clean_name.to_string();
45        let reeval_name = quote::format_ident!("{}_reevaluation_rule", clean_arg_name);
46        let input_span = input.span.clone();
47
48        match cache_strategy {
49            (CacheStrategy::None, _span) => quote_spanned! {input_span=>
50                // Non-transparent functions never clone, always move
51                let (#arg_name, #reeval_name): (#arg_type, directed::ReevaluationRule) = {
52                    match inputs.#arg_name.take() {
53                        Some((arg, reeval_name)) => (arg, reeval_name),
54                        None => return Err(directed::InjectionError::InputNotFound(Self::SHAPE.inputs.iter().find(|field| field.name == #clean_arg_name)))
55                    }
56                };
57            },
58            (CacheStrategy::Last, _span) | (CacheStrategy::All, _span) => quote_spanned! {input_span=>
59                let (#arg_name, #reeval_name): (#arg_type, directed::ReevaluationRule) = {
60                    match &inputs.#arg_name {
61                        Some((arg, reeval_name)) => (arg.clone(), *reeval_name),
62                        None => return Err(directed::InjectionError::InputNotFound(Self::SHAPE.inputs.iter().find(|field| field.name == #clean_arg_name)))
63                    }
64                };
65            },
66        }
67    }).collect())
68}
69
70/// This generates the code that uses the output of a parent node to set the
71/// input of a child node.
72fn input_injection(
73    stage_name: &syn::Ident,
74    inputs: &[InputParam],
75) -> Result<proc_macro2::TokenStream, Diagnostic> {
76    let mut match_arms = Vec::new();
77
78    // Build match arms from inputs
79    for input in inputs.iter() {
80        let clean_arg_name = input.clean_name.to_string();
81        let span = input.span.clone();
82
83        match_arms.push(quote_spanned! {span=>
84            Some(#clean_arg_name) => {
85                // Cast node to the concrete type
86                let node = node.as_any_mut()
87                    .downcast_mut::<directed::Node<Self>>()
88                    .ok_or_else(|| directed::InjectionError::InputTypeMismatch(input))?;
89                
90                // We need to handle different cases based on parent's reevaluation rule
91                if parent.reeval_rule() == directed::ReevaluationRule::Move {
92                    inject_move(node, parent, output, input)?;
93                } else {
94                    inject_clone(node, parent, output, input)?;
95                }
96                
97                Ok(())
98            }
99        });
100    }
101
102    // Add the default case
103    let default_case = quote_spanned! {Span::call_site()=>
104        Some(name) => {
105            // TODO: Verify the unwrap in unreachable
106            Err(directed::InjectionError::InputNotFound(#stage_name::SHAPE.inputs.iter().find(|field| field.name == name)))
107        },
108        None => Ok(()) // This means there's a connection with no data associated
109    };
110    match_arms.push(default_case);
111
112    // Generate the helper functions that work with concrete types
113    let inject_helpers = generate_inject_helpers(stage_name, &inputs)?;
114
115    Ok(quote_spanned! {Span::call_site()=>
116        #inject_helpers
117
118        let result: Result<(), directed::InjectionError> = match input.map(|input| input.name) {
119            #(#match_arms)*
120        };
121        result
122    })
123}
124
125fn generate_inject_helpers(
126    stage_name: &syn::Ident,
127    inputs: &[InputParam],
128) -> Result<proc_macro2::TokenStream, Diagnostic> {
129    let mut inject_move_arms = Vec::new();
130    let mut inject_clone_arms = Vec::new();
131
132    for input in inputs.iter() {
133        let arg_name = &input.clean_name;
134        let clean_arg_name = &input.clean_name.to_string();
135        let arg_type = true_type(&input.ty);
136        let span = input.span.clone();
137
138        inject_move_arms.push(quote_spanned! {span=>
139            Some(#clean_arg_name) => {
140                // For opaque parent, we need to remove the output
141                let output_val = parent.outputs_mut()
142                    .take_field(output.clone())
143                    .ok_or_else(|| directed::InjectionError::OutputNotFound(output.clone()))?;
144
145                let typed_val = output_val
146                    .downcast::<#arg_type>()
147                    .map_err(|_| directed::InjectionError::OutputTypeMismatch(output.clone()))?;
148
149                node.set_input_changed(true);
150
151                // TODO: Do something with old val
152                node.inputs.#arg_name.replace((*typed_val, directed::ReevaluationRule::Move));
153            }
154        });
155
156        inject_clone_arms.push(quote_spanned! {span=>
157            Some(#clean_arg_name) => {
158                // For transparent parent, clone the output
159                let output_val = parent.outputs_mut()
160                    .field_mut(output.clone())
161                    .ok_or_else(|| directed::InjectionError::OutputNotFound(output.clone()))?;
162
163                let typed_val = output_val
164                    .downcast_ref::<#arg_type>()
165                    .ok_or_else(|| directed::InjectionError::OutputTypeMismatch(output.clone()))?;
166                
167                // Check if changed
168                let input_changed = if let Some((existing_val, _)) = &node.inputs.#arg_name {
169                    *typed_val != *existing_val
170                } else {
171                    true
172                };
173                
174                if input_changed && !node.input_changed() {
175                    node.set_input_changed(true);
176                }
177                
178                // Set on concrete node
179                // TODO: Do something with old val
180                node.inputs.#arg_name.replace((typed_val.clone(), directed::ReevaluationRule::CacheLast));
181            }
182        });
183    }
184
185    Ok(quote_spanned! {Span::call_site()=>
186
187        #[allow(unreachable_code)]
188        fn inject_move(
189            node: &mut directed::Node<#stage_name>,
190            parent: &mut Box<dyn directed::AnyNode>,
191            output: Option<&'static directed::TypeReflection>,
192            input: Option<&'static directed::TypeReflection>
193        ) -> Result<(), directed::InjectionError> {
194            match input.map(|input| input.name) {
195                #(#inject_move_arms)*
196                _ => return Err(directed::InjectionError::InputNotFound(input))
197            }
198            Ok(())
199        }
200
201        #[allow(unreachable_code)]
202        fn inject_clone(
203            node: &mut directed::Node<#stage_name>,
204            parent: &mut Box<dyn directed::AnyNode>,
205            output: Option<&'static directed::TypeReflection>,
206            input: Option<&'static directed::TypeReflection>
207        ) -> Result<(), directed::InjectionError> {
208            match input.map(|input| input.name) {
209                #(#inject_clone_arms)*
210                _ => return Err(directed::InjectionError::InputNotFound(input))
211            }
212            Ok(())
213        }
214    })
215}
216
217fn generate_output_handling(
218    config: &StageConfig,
219    output_struct_name: &syn::Ident,
220    cache_strategy: (CacheStrategy, Span),
221) -> Result<proc_macro2::TokenStream, Diagnostic> {
222    let stage_name = &config.stage_name;
223    let arg_names = config
224        .inputs
225        .iter()
226        .map(|input| &input.clean_name)
227        .collect::<Vec<_>>();
228    let state_as_inputs = config
229        .states
230        .iter()
231        .map(|(name, _ty, span)| {
232            quote_spanned! {*span=> &mut state.#name}
233        })
234        .collect::<Vec<_>>();
235    let await_call = if config.is_async.0 {
236        quote::quote! {.await}
237    } else {
238        quote::quote! {}
239    };
240    let fn_call = match &config.outputs {
241        OutputParams::Explicit(_output_params) => {
242            quote_spanned! {Span::mixed_site()=> {
243                    #stage_name::call(#(#state_as_inputs),* #(#arg_names),*) #await_call
244                }
245            }
246        }
247        OutputParams::Implicit(_, span) => {
248            quote_spanned! {*span=>
249                #output_struct_name(Some(#stage_name::call(#(#state_as_inputs),* #(#arg_names),*) #await_call))
250            }
251        }
252    };
253
254    if cache_strategy.0 == CacheStrategy::All {
255        Ok(quote_spanned! {cache_strategy.1=>
256
257            // Use a hasher
258            let hash: u64 = {
259                #[allow(unused_imports)]
260                use std::hash::Hash;
261                #[allow(unused_imports)]
262                use std::hash::Hasher;
263                #[allow(unused_mut)]
264                let mut hasher = std::hash::DefaultHasher::new();
265                #(#arg_names.hash(&mut hasher);)*
266                hasher.finish()
267            };
268
269            // Check cache
270            #[allow(unused_variables)]
271            let cached = cache.get(&hash).and_then(|cached_vec| {
272                cached_vec.iter().find(|cached| {
273                    cached.inputs == *inputs
274                })
275            });
276
277            if let Some(cached) = cached {
278                // Just use cached values
279                Ok(cached.outputs.clone())
280            } else {
281                // Call and store result in cache
282                let result = #fn_call;
283
284                let cache_entry = directed::Cached {
285                    inputs: inputs.clone(),
286                    outputs: result.clone(),
287                };
288
289                cache.entry(hash).or_insert_with(Vec::new).push(cache_entry);
290                Ok(result)
291            }
292        })
293    } else {
294        // No advanced caching, just run it
295        Ok(quote_spanned!(cache_strategy.1=>Ok(#fn_call)))
296    }
297}
298
299fn prepare_input_types(config: &StageConfig) -> Result<Vec<proc_macro2::TokenStream>, Diagnostic> {
300    let args = config
301        .inputs
302        .iter()
303        .map(|input| (&input.clean_name, &input.ty, &input.ref_type));
304    let mut output = Vec::new();
305    for (arg_name, ty, ref_type) in args {
306        let reeval_name = quote::format_ident!("{}_reevaluation_rule", arg_name);
307        match ref_type {
308            RefType::Owned => {
309                output.push(quote_spanned!{arg_name.span()=>
310                    let #arg_name: #ty = match #reeval_name {
311                        directed::ReevaluationRule::Move => {
312                            // Parent is opaque
313                            #arg_name
314                        },
315                        directed::ReevaluationRule::CacheLast | directed::ReevaluationRule::CacheAll => {
316                            // Parent is transparent, clone the value
317                            #arg_name.clone()
318                        },
319                    };
320                });
321            }
322            RefType::Borrowed => {
323                output.push(quote_spanned! {arg_name.span()=>
324                    let #arg_name = &#arg_name;
325                });
326            }
327            RefType::BorrowedMut => {
328                output.push(quote_spanned! {arg_name.span()=>
329                    let #arg_name = &mut #arg_name;
330                });
331            },
332        }
333    }
334    Ok(output)
335}
336
337fn generate_dyn_fields_impl<P: Param>(params: impl Iterator<Item = P>) -> proc_macro2::TokenStream {
338    let mut field_arms: Vec<proc_macro2::TokenStream> = Vec::new();
339    let mut field_mut_arms: Vec<proc_macro2::TokenStream> = Vec::new();
340    let mut take_field_arms: Vec<proc_macro2::TokenStream> = Vec::new();
341    for param in params {
342        let (name, _, _span) = param.param();
343        match &name {
344            syn::Member::Named(ident) => {
345                let name_str = ident.to_string();
346                field_arms.push(quote_spanned! {Span::mixed_site()=>
347                    Some(#name_str) => self.#name.as_ref().map(|s| s as &dyn std::any::Any),
348                });
349                field_mut_arms.push(quote_spanned! {Span::mixed_site()=>
350                    Some(#name_str) => self.#name.as_mut().map(|s| s as &mut dyn std::any::Any),
351                });
352                take_field_arms.push(quote_spanned! {Span::mixed_site()=>
353                    Some(#name_str) => self.#name.take().map(|a| Box::new(a) as Box<dyn std::any::Any>),
354                });
355            }
356            syn::Member::Unnamed(index) => {
357                field_arms.push(quote_spanned! {Span::mixed_site()=>
358                    None => self.#index.as_ref().map(|s| s as &dyn std::any::Any),
359                });
360                field_mut_arms.push(quote_spanned! {Span::mixed_site()=>
361                    None => self.#index.as_mut().map(|s| s as &mut dyn std::any::Any),
362                });
363                take_field_arms.push(quote_spanned! {Span::mixed_site()=>
364                    None => self.#index.take().map(|a| Box::new(a) as Box<dyn std::any::Any>),
365                });
366            }
367        }
368    }
369    quote_spanned! {Span::mixed_site()=>
370       fn field<'a>(&'a self, field: Option<&'static directed::TypeReflection>) -> Option<&'a (dyn std::any::Any + 'static)> {
371            match field.map(|field| field.name) {
372                #(#field_arms)*
373                _ => None
374            }
375        }
376
377        fn field_mut<'a>(&'a mut self, field: Option<&'static directed::TypeReflection>) -> Option<&'a mut (dyn std::any::Any + 'static)> {
378            match field.map(|field| field.name) {
379                #(#field_mut_arms)*
380                _ => None
381            }
382        }
383
384        fn take_field(&mut self, field: Option<&'static directed::TypeReflection>) -> Option<Box<dyn std::any::Any>> {
385            match field.map(|field| field.name) {
386                #(#take_field_arms)*
387                _ => None
388            }
389        }
390    }
391}
392
393/// Modify the function to give it access to concrete output type
394fn insert_local_types(config: &mut StageConfig, output_type: proc_macro2::TokenStream) {
395    let original_fn = &mut config.original_fn;
396    let original_body = &mut original_fn.block;
397
398    original_body.stmts.insert(
399        0,
400        syn::Stmt::Item(syn::Item::Type(syn::ItemType {
401            attrs: Vec::new(),
402            vis: syn::Visibility::Inherited,
403            type_token: Token![type](Span::call_site()),
404            ident: syn::Ident::new("StageOutputType", Span::call_site()),
405            generics: syn::Generics::default(),
406            eq_token: Token![=](Span::call_site()),
407            ty: Box::new(syn::Type::Verbatim(output_type)),
408            semi_token: Token![;](Span::call_site()),
409        })),
410    );
411}
412
413/// The core trait that defines a stage - the culmination of this macro
414fn generate_stage_impl(mut config: StageConfig) -> Result<proc_macro2::TokenStream, Diagnostic> {
415    // Generate names for types
416    let stage_name_str = config.stage_name.to_string();
417    let input_struct_name = quote::format_ident! {"{}InputCache", &config.stage_name};
418    let output_struct_name = quote::format_ident! {"{}OutputCache", &config.stage_name};
419    let state_struct_name = quote::format_ident! {"{}State", &config.stage_name};
420
421    // Modify stage before starting the generation
422    insert_local_types(&mut config, quote::quote! {#output_struct_name});
423
424    let original_fn = &config.original_fn;
425    let stage_name = &config.stage_name;
426    let states = &config.states;
427    let fn_attrs = &original_fn.attrs;
428    let fn_vis = &original_fn.vis;
429    let original_args = &original_fn.sig.inputs;
430    let fn_return_type = &original_fn.sig.output;
431    let original_body = &original_fn.block;
432    let async_status = &original_fn.sig.asyncness;
433
434    // Determine derives for inputs
435    let input_derives = match config.cache_strategy.0 {
436        CacheStrategy::None => {
437            quote_spanned! {config.cache_strategy.1=>#[derive(Default)]}
438        }
439        CacheStrategy::Last => {
440            quote_spanned! {config.cache_strategy.1=>#[derive(Default, Clone, PartialEq)]}
441        }
442        CacheStrategy::All => {
443            quote_spanned! {config.cache_strategy.1=>#[derive(Default, Clone, PartialEq, Eq, Hash)]}
444        }
445    };
446    // Determine derives for outputs
447    let output_derives = match config.cache_strategy.0 {
448        CacheStrategy::None => quote::quote! {#[derive(Default)]},
449        CacheStrategy::Last | CacheStrategy::All => {
450            quote_spanned! {config.cache_strategy.1=>#[derive(Default, Clone)]}
451        }
452    };
453
454    // Create useful structs
455    let input_struct_fields =
456        proc_macro2::TokenStream::from_iter(config.inputs.iter().map(|input| {
457            let input_ident = &input.clean_name;
458            let input_ty = input.unwrapped_type();
459            quote_spanned! {input.span=>
460                #input_ident: Option<(#input_ty, directed::ReevaluationRule)>,
461            }
462        }));
463    let input_struct = quote_spanned! {Span::call_site()=>
464        #input_derives
465        #fn_vis struct #input_struct_name {
466            #input_struct_fields
467        }
468    };
469    let input_dyn_fields_trait_impl = generate_dyn_fields_impl(config.inputs.iter());
470    let input_shape = {
471        let input_fields = config.inputs.iter().map(|field| {
472            let name = field.clean_name.to_string();
473            let ty = &field.ty;
474            let ty_string = quote::quote!{#ty}.to_string();
475            quote::quote!{
476                directed::TypeReflection { name: #name, ty: #ty_string }
477            }
478        });
479        quote::quote!{
480            impl #input_struct_name {
481                const SHAPE: &'static [directed::TypeReflection] = &[#(#input_fields),*];
482            }
483        }
484    };
485
486    let output_struct = match &config.outputs {
487        OutputParams::Explicit(output_params) => {
488            let mut fields = Vec::new();
489            for param in output_params.iter() {
490                let name = &param.name;
491                let ty = &param.ty;
492                let span = param.span;
493
494                fields.push(quote_spanned! {span=> #name: Option<#ty>});
495            }
496            quote_spanned! {Span::mixed_site()=>
497                #output_derives
498                #fn_vis struct #output_struct_name {
499                    #(#fields),*
500                }
501            }
502        }
503        OutputParams::Implicit(ty, span) => quote_spanned! {*span=>
504            #output_derives
505            #fn_vis struct #output_struct_name(Option<#ty>);
506        },
507    };
508
509    let output_dyn_fields_trait_impl = match &config.outputs {
510        OutputParams::Explicit(output_params) => generate_dyn_fields_impl(output_params.iter()),
511        OutputParams::Implicit(_, span) => generate_dyn_fields_impl(std::iter::once(TupleParam {
512            idx: syn::Index::from(0),
513            ty: syn::TypeNever {
514                bang_token: syn::Token![!](*span),
515            }
516            .into(),
517            span: Span::mixed_site(),
518        })),
519    };
520    let output_shape = {
521        let output_fields = match &config.outputs {
522            OutputParams::Explicit(output_params) => output_params.iter().map(|field| {
523                let name = field.name.to_string();
524                let ty = &field.ty;
525                let ty_string = quote::quote!{#ty}.to_string();
526                quote::quote_spanned!{field.span=>
527                    directed::TypeReflection { name: #name, ty: #ty_string }
528                }
529            }).collect(),
530            OutputParams::Implicit(ty, span) => {
531                let ty_string = quote::quote!{#ty}.to_string();
532                vec!(quote::quote_spanned!{*span=>
533                    directed::TypeReflection { name: "_", ty: #ty_string }
534                })
535            },
536        };
537        quote::quote!{
538            impl #output_struct_name {
539                const SHAPE: &'static [directed::TypeReflection] = &[#(#output_fields),*];
540            }
541        }
542    };
543
544    let state_struct_fields =
545        proc_macro2::TokenStream::from_iter(states.iter().map(|(state_ident, ty, span)| {
546            let state_ty = &ty;
547            quote_spanned! {*span=>
548                #state_ident: #state_ty,
549            }
550        }));
551    // If state is a unit, implement default
552    let default_derive = if config.states.is_empty() {
553        quote_spanned! { Span::call_site()=> #[derive(Default)] }
554    } else {
555        quote::quote!{}
556    };
557    let state_struct = quote_spanned! {Span::call_site()=>
558        #default_derive
559        #fn_vis struct #state_struct_name {
560            #state_struct_fields
561        }
562    };
563
564    // Generate code sections
565    let extraction_code = generate_extraction_code(&config.inputs, config.cache_strategy)?;
566    let injection_code = input_injection(stage_name, &config.inputs)?;
567    let prepare_input_types_code = prepare_input_types(&config)?;
568    let output_handling =
569        generate_output_handling(&config, &output_struct_name, config.cache_strategy)?;
570
571    // Determine evaluation strategy and reevaluation rule
572    let eval_strategy = if config.is_lazy.0 {
573        quote_spanned! {config.is_lazy.1=> directed::EvalStrategy::Lazy }
574    } else {
575        quote_spanned! {config.is_lazy.1=> directed::EvalStrategy::Urgent }
576    };
577
578    let reevaluation_rule = match &config.cache_strategy {
579        (CacheStrategy::None, span) => quote_spanned! {*span=> directed::ReevaluationRule::Move },
580        (CacheStrategy::Last, span) => {
581            quote_spanned! {*span=> directed::ReevaluationRule::CacheLast }
582        }
583        (CacheStrategy::All, span) => {
584            quote_spanned! {*span=> directed::ReevaluationRule::CacheAll }
585        }
586    };
587
588    // Generate function inputs
589    let state_as_input = states
590        .iter()
591        .map(|(name, ty, span)| quote_spanned! {*span=>#name: &mut #ty});
592
593    // Redefine the original function
594    let return_type = match &config.outputs {
595        OutputParams::Explicit(_) => quote::quote! {-> #output_struct_name},
596        OutputParams::Implicit(_, span) => quote::quote_spanned! {*span => #fn_return_type},
597    };
598    let call_fn_def = quote::quote! {
599        #(#fn_attrs)*
600        #async_status fn call(#(#state_as_input,)* #original_args) #return_type #original_body
601    };
602
603    // Slap together eval logic
604    let eval_logic = quote_spanned! {Span::call_site()=>
605        #(#extraction_code)*
606        #(#prepare_input_types_code)*
607        #output_handling
608    };
609    let sync_eval_logic = if async_status.is_some() {
610        quote::quote!(panic!("Attempted to call async code synchronously"))
611    } else {
612        quote::quote!(#eval_logic)
613    };
614    let async_eval_logic = if async_status.is_some() {
615        quote::quote!(#eval_logic)
616    } else {
617        quote::quote!(#eval_logic)
618    };
619
620    let async_trait_derive = if cfg!(feature = "tokio") {
621        quote::quote!{#[cfg_attr(feature = "tokio", async_trait::async_trait)]}
622    } else {
623        quote::quote!{}
624    };
625
626    let evaluate_impls = if cfg!(feature = "tokio") {
627        quote::quote!{
628            #[cfg(feature = "tokio")]
629            async fn evaluate_async(
630                &self,
631                state: &mut Self::State,
632                inputs: &mut Self::Input,
633                cache: &mut std::collections::HashMap<u64, Vec<directed::Cached<Self>>>,
634            ) -> Result<Self::Output, InjectionError> {
635                #async_eval_logic
636            }
637
638            fn evaluate(
639                &self,
640                state: &mut Self::State,
641                inputs: &mut Self::Input,
642                cache: &mut std::collections::HashMap<u64, Vec<directed::Cached<Self>>>,
643            ) -> Result<Self::Output, InjectionError> {
644                #sync_eval_logic
645            }
646        }
647    } else {
648        quote::quote!{
649            fn evaluate(
650                &self,
651                state: &mut Self::State,
652                inputs: &mut Self::Input,
653                cache: &mut std::collections::HashMap<u64, Vec<directed::Cached<Self>>>,
654            ) -> Result<Self::Output, InjectionError> {
655                #sync_eval_logic
656            }
657        }
658    };
659
660    // The coup de grace
661    Ok(quote_spanned! {Span::call_site()=>
662        // Create a struct implementing the Stage trait
663        #[derive(Debug, Clone, Copy, Default)]
664        #fn_vis struct #stage_name;
665
666        impl #stage_name {
667            /// Evaluate this stage
668            #call_fn_def
669        }
670
671        // Structs to contain various caches
672        #input_struct
673        #input_shape
674        #output_struct
675        #output_shape
676        #state_struct
677
678        impl directed::DynFields for #output_struct_name {
679            #output_dyn_fields_trait_impl
680
681            fn replace(&mut self, other: Box<dyn std::any::Any>) -> Box<dyn DynFields> {
682                Box::new(std::mem::replace::<Self>(self, *other.downcast::<Self>().expect("DynFields type must be exact match")))
683            }
684        }
685
686        impl directed::DynFields for #input_struct_name {
687            #input_dyn_fields_trait_impl
688
689            fn replace(&mut self, other: Box<dyn std::any::Any>) -> Box<dyn DynFields> {
690                Box::new(std::mem::replace::<Self>(self, *other.downcast::<Self>().expect("DynFields type must be exact match")))
691            }
692        }
693
694        #async_trait_derive
695        impl directed::Stage for #stage_name {
696            const SHAPE: directed::StageShape = directed::StageShape {
697                stage_name: #stage_name_str,
698                inputs: Self::Input::SHAPE,
699                outputs: Self::Output::SHAPE,
700            };
701            type Input = #input_struct_name;
702            type Output = #output_struct_name;
703            type State = #state_struct_name;
704
705            #evaluate_impls
706
707            fn eval_strategy(&self) -> directed::EvalStrategy {
708                #eval_strategy
709            }
710
711            fn reeval_rule(&self) -> directed::ReevaluationRule {
712                #reevaluation_rule
713            }
714
715            fn inject_input(&self, node: &mut directed::Node<Self>, parent: &mut Box<dyn directed::AnyNode>, output: Option<&'static directed::TypeReflection>, input: Option<&'static directed::TypeReflection>) -> Result<(), directed::InjectionError> {
716                #injection_code
717            }
718        }
719    })
720}