trident_syn/parser/
trident_flow_executor.rs

1use proc_macro2::TokenStream;
2use syn::parse::Error as ParseError;
3use syn::parse::Result as ParseResult;
4use syn::parse::{Parse, ParseStream};
5use syn::spanned::Spanned;
6use syn::ItemImpl;
7use syn::Meta;
8
9use crate::types::trident_flow_executor::FlowExecutorArgs;
10use crate::types::trident_flow_executor::TridentFlowExecutorImpl;
11
12impl Parse for FlowExecutorArgs {
13    fn parse(input: ParseStream) -> ParseResult<Self> {
14        let mut args = FlowExecutorArgs::default();
15
16        while !input.is_empty() {
17            let meta: Meta = input.parse()?;
18
19            match meta {
20                Meta::NameValue(nv) => {
21                    if nv.path.is_ident("random_tail") {
22                        if let syn::Expr::Lit(expr_lit) = nv.value {
23                            if let syn::Lit::Bool(lit_bool) = expr_lit.lit {
24                                args.random_tail = lit_bool.value();
25                            } else {
26                                return Err(ParseError::new(
27                                    expr_lit.lit.span(),
28                                    "random_tail must be a boolean value",
29                                ));
30                            }
31                        }
32                    } else {
33                        return Err(ParseError::new(
34                            nv.path.span(),
35                            format!("unknown attribute: {}", nv.path.get_ident().unwrap()).as_str(),
36                        ));
37                    }
38                }
39                Meta::Path(path) => {
40                    return Err(ParseError::new(
41                        path.span(),
42                        format!("unknown flag attribute: {}", path.get_ident().unwrap()).as_str(),
43                    ));
44                    // if path.is_ident("shuffle") {
45                    //     args.shuffle = true;
46                    // } else {
47                    //     return Err(ParseError::new(
48                    //         path.span(),
49                    //         format!("unknown flag attribute: {}", path.get_ident().unwrap())
50                    //             .as_str(),
51                    //     ));
52                    // }
53                }
54                _ => {
55                    return Err(ParseError::new(
56                        meta.span(),
57                        "expected either a name-value pair or a flag attribute",
58                    ));
59                }
60            }
61
62            // Parse comma if there are more attributes
63            if !input.is_empty() {
64                input.parse::<syn::Token![,]>()?;
65            }
66        }
67
68        Ok(args)
69    }
70}
71
72pub fn parse_trident_flow_executor(
73    attr: TokenStream,
74    input: &ItemImpl,
75) -> ParseResult<TridentFlowExecutorImpl> {
76    let args: FlowExecutorArgs = syn::parse2(attr)?;
77
78    // Extract just the path without any generics
79    let type_name = if let syn::Type::Path(type_path) = &*input.self_ty {
80        let mut cleaned_path = type_path.clone();
81        // Clear any generic arguments from the last segment
82        if let Some(last_segment) = cleaned_path.path.segments.last_mut() {
83            last_segment.arguments = syn::PathArguments::None;
84        }
85        Box::new(syn::Type::Path(cleaned_path))
86    } else {
87        input.self_ty.clone()
88    };
89    let generics = input.generics.clone();
90
91    let mut init_method = None;
92    let mut flow_methods = Vec::new();
93
94    // Collect init and flow methods
95    for item in &input.items {
96        if let syn::ImplItem::Fn(method) = item {
97            // First check for init methods
98            if method.attrs.iter().any(|attr| attr.path().is_ident("init")) {
99                if init_method.is_some() {
100                    return Err(ParseError::new(
101                        method.span(),
102                        "Multiple #[init] methods found. Only one is allowed.",
103                    ));
104                }
105                init_method = Some(method.sig.ident.clone());
106                continue;
107            }
108
109            // Then check for flow methods
110            if method.attrs.iter().any(|attr| attr.path().is_ident("flow")) {
111                // Only check for ignore if it's a flow method
112                let is_ignored = method
113                    .attrs
114                    .iter()
115                    .any(|attr| attr.path().is_ident("flow_ignore"));
116                if !is_ignored {
117                    flow_methods.push(method.sig.ident.clone());
118                }
119            }
120        }
121    }
122
123    Ok(TridentFlowExecutorImpl {
124        type_name,
125        impl_block: input.items.clone(),
126        flow_methods,
127        init_method,
128        generics,
129        args,
130    })
131}