1use proc_macro::{TokenStream, Span};
2use quote::quote;
3use std::collections::HashMap;
4use syn::parse::{Parse, ParseStream, Result};
5use syn::{
6 braced, parse_macro_input, parse_quote, punctuated::Punctuated, Arm, Expr, Ident, LitChar,
7 LitStr, Token,
8};
9
10struct BuildTrie {
11 state_enum_name: Ident,
12 result_enum_name: Ident,
13 function_name: Ident,
14 result_name: Ident,
15 mappings: Punctuated<(LitStr, Expr), Token![,]>,
16}
17
18impl Parse for BuildTrie {
19 fn parse(input: ParseStream) -> Result<Self> {
20 let mut function_name: Option<Ident> = None;
21 let mut state_enum_name: Option<Ident> = None;
22 let mut result_enum_name: Option<Ident> = None;
23 let mut result_name: Option<Ident> = None;
24 let mut mappings: Option<Punctuated<(LitStr, Expr), Token![,]>> = None;
25 while !(state_enum_name.is_some()
26 && result_enum_name.is_some()
27 && function_name.is_some()
28 && mappings.is_some()
29 && result_name.is_some())
30 && !input.is_empty()
31 {
32 let key: Ident = input.parse()?;
33 input.parse::<Token![:]>()?;
34 match key.to_string().as_str() {
35 "function" => {
36 if function_name.is_some() {
37 panic!("Function name already defined ")
38 }
39 input.parse::<Token![fn]>()?;
40 let name: Ident = input.parse()?;
41 input.parse::<Token![;]>()?;
42 function_name = Some(name);
43 }
44 "state_enum" => {
45 if state_enum_name.is_some() {
46 panic!("State enum name already defined ")
47 }
48 input.parse::<Token![enum]>()?;
49 let name: Ident = input.parse()?;
50 input.parse::<Token![;]>()?;
51 state_enum_name = Some(name);
52 }
53 "result_enum" => {
54 if result_enum_name.is_some() {
55 panic!("Result enum name already defined ")
56 }
57 input.parse::<Token![enum]>()?;
58 let name: Ident = input.parse()?;
59 input.parse::<Token![;]>()?;
60 result_enum_name = Some(name);
61 }
62 "result" => {
63 if result_name.is_some() {
64 panic!("Result reference already defined ")
65 }
66 let name: Ident = input.parse()?;
67 input.parse::<Token![;]>()?;
68 result_name = Some(name);
69 }
70 "mappings" => {
71 if mappings.is_some() {
72 panic!("Mappings already defined ")
73 }
74 let content;
75 braced!(content in input);
76 let mappings_result =
77 content.parse_terminated::<(LitStr, Expr), Token![,]>(|input| {
78 let string: LitStr = input.parse()?;
79 input.parse::<Token![=>]>()?;
80 let expr: Expr = input.parse()?;
81 Ok((string, expr))
82 })?;
83 mappings = Some(mappings_result);
84 }
85 _ => panic!("invalid key"),
86 }
87 }
88
89 Ok(BuildTrie {
90 state_enum_name: state_enum_name.expect("No state enum name"),
91 result_enum_name: result_enum_name.expect("No result enum name"),
92 function_name: function_name.expect("No function name"),
93 mappings: mappings.expect("No mappings"),
94 result_name: result_name.expect("No result name"),
95 })
96 }
97}
98
99struct Trie<K, V>(HashMap<K, Trie<K, V>>, Option<V>);
100
101impl<K, V> Trie<K, V> {
102 fn is_leaf(&self) -> bool {
103 self.0.is_empty()
104 }
105}
106
107const NO_STATE_NAME: &str = "None";
108
109#[proc_macro]
110pub fn build_trie(input: TokenStream) -> TokenStream {
111 let BuildTrie {
112 state_enum_name,
113 result_enum_name,
114 function_name,
115 result_name,
116 mappings,
117 } = parse_macro_input!(input as BuildTrie);
118
119 let mut trie: Trie<char, Expr> = Trie(HashMap::new(), None);
120
121 for (string, value) in mappings {
122 let mut node = &mut trie;
123 for chr in string.value().chars() {
124 if node.0.get(&chr).is_none() {
125 node.0.insert(chr, Trie(HashMap::new(), None));
126 }
127 node = node.0.get_mut(&chr).unwrap();
128 }
129 node.1 = Some(value);
130 }
131
132 let mut states: Vec<Ident> = Vec::new();
133 let mut arms: Vec<Arm> = Vec::new();
134
135 fn expand_trie(
136 trie: &Trie<char, Expr>,
137 state_enum_name_ident: &Ident,
138 result_enum_name_ident: &Ident,
139 arms: &mut Vec<Arm>,
140 states: &mut Vec<Ident>,
141 prev_state: &Ident
142 ) {
143 let mut count: u8 = 0;
144 for (key, value) in trie.0.iter() {
145 let chr = LitChar::new(*key, Span::call_site().into());
146 if value.is_leaf() {
147 if let Some(value) = &value.1 {
148 let arm: Arm = parse_quote! {
149 (#state_enum_name_ident::#prev_state, #chr) => #result_enum_name_ident::Result(#value, true),
150 };
151 arms.push(arm);
152 }
153 } else {
154 let new_state = {
155 let as_string = prev_state.to_string();
156 count += 1;
157 if as_string.is_empty() || as_string == NO_STATE_NAME {
158 let mut string = String::new();
159 string.push((count + 96) as char);
160 Ident::new(&string, Span::call_site().into())
161 } else {
162 let mut string = as_string.clone();
163 string.push((count + 96) as char);
164 Ident::new(&string, Span::call_site().into())
165 }
166 };
167 states.push(new_state.clone());
168 let arm: Arm = parse_quote! {
169 (#state_enum_name_ident::#prev_state, #chr) => #result_enum_name_ident::NewState(#state_enum_name_ident::#new_state),
170 };
171 arms.push(arm);
172 expand_trie(
173 value,
174 state_enum_name_ident,
175 result_enum_name_ident,
176 arms,
177 states,
178 &new_state
179 );
180 if let Some(value) = &value.1 {
181 let arm: Arm = parse_quote! {
182 (#state_enum_name_ident::#new_state, _) => #result_enum_name_ident::Result(#value, false),
183 };
184 arms.push(arm);
185 }
186 }
187 }
188 }
189
190 let no_state_ident = Ident::new(NO_STATE_NAME, Span::call_site().into());
191
192 expand_trie(
193 &trie,
194 &state_enum_name,
195 &result_enum_name,
196 &mut arms,
197 &mut states,
198 &no_state_ident
199 );
200
201 let expanded = quote! {
202 pub enum #result_enum_name {
203 Result(#result_name, bool),
204 NewState(#state_enum_name)
205 }
206
207 #[derive(PartialEq, Debug, Copy, Clone)]
208 pub enum #state_enum_name {
209 None,
210 #( #states ),*
211 }
212
213 pub fn #function_name(state: &#state_enum_name, chr: &char) -> #result_enum_name {
214 match (state, chr) {
215 #( #arms )*
216 (#state_enum_name::#no_state_ident, _) => #result_enum_name::NewState(#state_enum_name::#no_state_ident),
217 (state, chr) => panic!("Invalid {}", chr)
218 }
219 }
220 };
221
222 TokenStream::from(expanded)
223}