1#![recursion_limit = "128"]
5extern crate proc_macro;
6
7use proc_macro::TokenStream;
8use quote::{quote, ToTokens};
9use std::{collections::BTreeSet, iter::FromIterator};
10use syn::{parse_macro_input, Attribute, Ident};
11
12mod parser;
13
14struct Transition<'a> {
17 initial_state: &'a Ident,
18 input_value: &'a Ident,
19 final_state: &'a Ident,
20 output: &'a Option<Ident>,
21}
22
23fn attrs_to_token_stream(attrs: Vec<Attribute>) -> proc_macro2::TokenStream {
24 let attrs = attrs.into_iter().map(ToTokens::into_token_stream);
25 proc_macro2::TokenStream::from_iter(attrs)
26}
27
28#[proc_macro]
29pub fn state_machine(tokens: TokenStream) -> TokenStream {
32 let input = parse_macro_input!(tokens as parser::StateMachineDef);
33
34 let doc = attrs_to_token_stream(input.doc);
35 let attrs = attrs_to_token_stream(input.attributes);
36
37 if input.transitions.is_empty() {
38 let output = quote! {
39 compile_error!("rust-fsm: at least one state transition must be provided");
40 };
41 return output.into();
42 }
43
44 let fsm_name = input.name;
45 let visibility = input.visibility;
46
47 let transitions = input.transitions.iter().flat_map(|def| {
48 def.transitions.iter().map(move |transition| Transition {
49 initial_state: &def.initial_state,
50 input_value: &transition.input_value,
51 final_state: &transition.final_state,
52 output: &transition.output,
53 })
54 });
55
56 let mut states = BTreeSet::new();
57 let mut inputs = BTreeSet::new();
58 let mut outputs = BTreeSet::new();
59 let mut transition_cases = Vec::new();
60 let mut output_cases = Vec::new();
61
62 #[cfg(feature = "diagram")]
63 let mut mermaid_diagram = format!(
64 "///```mermaid\n///stateDiagram-v2\n/// [*] --> {}\n",
65 input.initial_state
66 );
67
68 states.insert(&input.initial_state);
69
70 for transition in transitions {
71 let Transition {
72 initial_state,
73 final_state,
74 input_value,
75 output,
76 } = transition;
77
78 #[cfg(feature = "diagram")]
79 mermaid_diagram.push_str(&format!(
80 "/// {initial_state} --> {final_state}: {input_value}"
81 ));
82
83 transition_cases.push(quote! {
84 (Self::State::#initial_state, Self::Input::#input_value) => {
85 Some(Self::State::#final_state)
86 }
87 });
88
89 if let Some(output_value) = output {
90 output_cases.push(quote! {
91 (Self::State::#initial_state, Self::Input::#input_value) => {
92 Some(Self::Output::#output_value)
93 }
94 });
95
96 #[cfg(feature = "diagram")]
97 mermaid_diagram.push_str(&format!(" [{output_value}]"));
98 }
99
100 #[cfg(feature = "diagram")]
101 mermaid_diagram.push('\n');
102
103 states.insert(initial_state);
104 states.insert(final_state);
105 inputs.insert(input_value);
106 if let Some(ref output) = output {
107 outputs.insert(output);
108 }
109 }
110
111 #[cfg(feature = "diagram")]
112 mermaid_diagram.push_str("///```");
113 #[cfg(feature = "diagram")]
114 let mermaid_diagram: proc_macro2::TokenStream = mermaid_diagram.parse().unwrap();
115
116 let initial_state_name = &input.initial_state;
117
118 let (input_type, input_impl) = match input.input_type {
119 Some(t) => (quote!(#t), quote!()),
120 None => (
121 quote!(Input),
122 quote! {
123 #attrs
124 pub enum Input {
125 #(#inputs),*
126 }
127 },
128 ),
129 };
130
131 let (state_type, state_impl) = match input.state_type {
132 Some(t) => (quote!(#t), quote!()),
133 None => (
134 quote!(State),
135 quote! {
136 #attrs
137 pub enum State {
138 #(#states),*
139 }
140 },
141 ),
142 };
143
144 let (output_type, output_impl) = match input.output_type {
145 Some(t) => (quote!(#t), quote!()),
146 None => {
147 let attrs = if outputs.is_empty() {
150 quote!()
151 } else {
152 attrs.clone()
153 };
154 (
155 quote!(Output),
156 quote! {
157 #attrs
158 pub enum Output {
159 #(#outputs),*
160 }
161 },
162 )
163 }
164 };
165
166 #[cfg(feature = "diagram")]
167 let diagram = quote! {
168 #[cfg_attr(doc, ::rust_fsm::aquamarine)]
169 #mermaid_diagram
170 };
171
172 #[cfg(not(feature = "diagram"))]
173 let diagram = quote!();
174
175 let output = quote! {
176 #doc
177 #diagram
178 #visibility mod #fsm_name {
179 #attrs
180 pub struct Impl;
181
182 pub type StateMachine = ::rust_fsm::StateMachine<Impl>;
183
184 #input_impl
185 #state_impl
186 #output_impl
187
188 impl ::rust_fsm::StateMachineImpl for Impl {
189 type Input = #input_type;
190 type State = #state_type;
191 type Output = #output_type;
192 const INITIAL_STATE: Self::State = Self::State::#initial_state_name;
193
194 fn transition(state: &Self::State, input: &Self::Input) -> Option<Self::State> {
195 match (state, input) {
196 #(#transition_cases)*
197 _ => None,
198 }
199 }
200
201 fn output(state: &Self::State, input: &Self::Input) -> Option<Self::Output> {
202 match (state, input) {
203 #(#output_cases)*
204 _ => None,
205 }
206 }
207 }
208 }
209 };
210
211 output.into()
212}