be_generust/
lib.rs

1mod utils;
2mod walker;
3
4use proc_macro::TokenStream;
5use quote::quote;
6use syn::spanned::Spanned;
7use syn::{parse_macro_input, ItemFn};
8
9use utils::{get_iter_item_type, make_ident, to_pascal_case};
10use walker::Walker;
11
12#[proc_macro_attribute]
13pub fn giver(attr: TokenStream, item: TokenStream) -> TokenStream {
14    let func = parse_macro_input!(item as ItemFn);
15
16    let iter_item_type = match get_iter_item_type(&func.sig.output) {
17        Some(ty) => ty,
18        None => {
19            return fail(
20                &func.sig.output,
21                "return type must be `-> impl Iterator<Item = XXX>`",
22            )
23        }
24    };
25
26    let func_vis = &func.vis;
27
28    let name_snake = func.sig.ident.to_string();
29    let name_pascal = to_pascal_case(&name_snake);
30
31    let func_name = make_ident(&name_snake);
32    let mod_name = make_ident(&format!("{}_mod", name_snake));
33    let state_enum_name = make_ident(&format!("{}State", name_pascal));
34    let struct_name = make_ident(&name_pascal);
35
36    let w = Walker::walk(state_enum_name.clone(), &func.block.stmts);
37    let state_idents = &w.states;
38    let match_blocks = w.output.iter().map(|((_, s), b)| {
39        let state_enum = &w.name;
40
41        if b.is_empty() {
42            quote! {
43                #state_enum::#s |
44            }
45        } else {
46            quote! {
47                #state_enum::#s => {
48                    #(#b)*
49                },
50            }
51        }
52    });
53
54    let new_code = quote! {
55        mod #mod_name {
56            enum #state_enum_name { #(#state_idents),* }
57
58            struct #struct_name {
59                state: #state_enum_name,
60            }
61
62            impl Iterator for #struct_name {
63                type Item = #iter_item_type;
64
65                fn next(&mut self) -> Option<#iter_item_type> {
66                    loop {
67                        match self.state {
68                            #(#match_blocks)*
69                        }
70                    }
71                }
72            }
73
74            pub fn #func_name() -> impl Iterator<Item = #iter_item_type> {
75                #struct_name { state: #state_enum_name::S0_Start }
76            }
77        }
78
79        #func_vis use #mod_name::#func_name;
80    };
81
82    if attr.to_string() == "print" {
83        println!("{}", &new_code);
84    }
85
86    return TokenStream::from(new_code);
87}
88
89fn fail<T: Spanned>(s: &T, msg: &str) -> TokenStream {
90    let msg = format!("[generoust] {}", msg);
91    let err = syn::Error::new(s.span(), msg).to_compile_error();
92
93    return TokenStream::from(err);
94}