Skip to main content

statum_macros/
lib.rs

1#![feature(proc_macro_span)]
2moddef::moddef!(
3    flat (pub) mod {
4    },
5    flat (pub(crate)) mod {
6        source_cache,
7        state,
8        machine,
9        transition,
10        validators
11    }
12);
13
14use crate::{
15    MachinePath, StateModulePath, ensure_machine_loaded, ensure_state_enum_loaded, get_machine,
16};
17use proc_macro::TokenStream;
18use syn::{ItemEnum, ItemImpl, ItemStruct, parse_macro_input};
19
20use module_path_extractor::get_pseudo_module_path;
21
22#[proc_macro_attribute]
23pub fn state(_attr: TokenStream, item: TokenStream) -> TokenStream {
24    let input = parse_macro_input!(item as ItemEnum);
25
26    // Validate the enum before proceeding
27    if let Some(error) = validate_state_enum(&input) {
28        return error.into();
29    }
30
31    let enum_info = match EnumInfo::from_item_enum(&input) {
32        Ok(info) => info,
33        Err(err) => return err.to_compile_error().into(),
34    };
35
36    // Store metadata in `state_enum_map`
37    store_state_enum(&enum_info);
38
39    // Generate structs and implementations dynamically
40    let expanded = generate_state_impls(&enum_info.module_path);
41
42    TokenStream::from(expanded)
43}
44
45#[proc_macro_attribute]
46pub fn machine(_attr: TokenStream, item: TokenStream) -> TokenStream {
47    let input = parse_macro_input!(item as ItemStruct);
48
49    let machine_info = MachineInfo::from_item_struct(&input);
50
51    // Validate the struct before proceeding
52    if let Some(error) = validate_machine_struct(&input, &machine_info) {
53        return error.into();
54    }
55
56    // Store metadata in `machine_map`
57    store_machine_struct(&machine_info);
58
59    // Generate any required structs or implementations dynamically
60    let expanded = generate_machine_impls(&machine_info);
61
62    TokenStream::from(expanded)
63}
64
65#[proc_macro_attribute]
66pub fn transition(
67    _attr: proc_macro::TokenStream,
68    item: proc_macro::TokenStream,
69) -> proc_macro::TokenStream {
70    let input = parse_macro_input!(item as ItemImpl);
71
72    // -- Step 1: Parse
73    let tr_impl = parse_transition_impl(&input);
74
75    let module_path = get_pseudo_module_path();
76
77    let state_path: StateModulePath = module_path.clone().into();
78    let machine_path: MachinePath = module_path.clone().into();
79    let _ = ensure_state_enum_loaded(&state_path);
80    let _ = ensure_machine_loaded(&machine_path);
81
82    let machine_info_owned = get_machine(&machine_path);
83    let machine_info = machine_info_owned
84        .as_ref()
85        .expect("Machine not found, even though we validated above");
86
87    if let Some(err) = validate_transition_functions(&tr_impl.functions, machine_info) {
88        return err.into();
89    }
90
91    // -- Step 3: Generate new code
92    let expanded = generate_transition_impl(&input, &tr_impl, machine_info, &module_path);
93
94    // Combine expanded code with the original `impl` if needed
95    // or simply return the expanded code
96    expanded.into()
97}
98
99#[proc_macro_attribute]
100pub fn validators(attr: TokenStream, item: TokenStream) -> TokenStream {
101    let module_path = get_pseudo_module_path();
102    parse_validators(attr, item, &module_path)
103}