1use phf_generator::HashState;
7use phf_shared::PhfHash;
8use proc_macro::TokenStream;
9use quote::quote;
10use std::collections::HashSet;
11use std::hash::Hasher;
12use syn::parse::{self, Parse, ParseStream};
13use syn::punctuated::Punctuated;
14use syn::{parse_macro_input, Error, Expr, ExprLit, Lit, Token, UnOp};
15#[cfg(feature = "unicase")]
16use unicase_::UniCase;
17
18#[derive(Hash, PartialEq, Eq, Clone)]
19enum ParsedKey {
20 Str(String),
21 Binary(Vec<u8>),
22 Char(char),
23 I8(i8),
24 I16(i16),
25 I32(i32),
26 I64(i64),
27 I128(i128),
28 Isize(isize),
29 U8(u8),
30 U16(u16),
31 U32(u32),
32 U64(u64),
33 U128(u128),
34 Usize(usize),
35 Bool(bool),
36 #[cfg(feature = "unicase")]
37 UniCase(UniCase<String>),
38}
39
40impl PhfHash for ParsedKey {
41 fn phf_hash<H>(&self, state: &mut H)
42 where
43 H: Hasher,
44 {
45 match self {
46 ParsedKey::Str(s) => s.phf_hash(state),
47 ParsedKey::Binary(s) => s.phf_hash(state),
48 ParsedKey::Char(s) => s.phf_hash(state),
49 ParsedKey::I8(s) => s.phf_hash(state),
50 ParsedKey::I16(s) => s.phf_hash(state),
51 ParsedKey::I32(s) => s.phf_hash(state),
52 ParsedKey::I64(s) => s.phf_hash(state),
53 ParsedKey::I128(s) => s.phf_hash(state),
54 ParsedKey::Isize(s) => s.phf_hash(state),
55 ParsedKey::U8(s) => s.phf_hash(state),
56 ParsedKey::U16(s) => s.phf_hash(state),
57 ParsedKey::U32(s) => s.phf_hash(state),
58 ParsedKey::U64(s) => s.phf_hash(state),
59 ParsedKey::U128(s) => s.phf_hash(state),
60 ParsedKey::Usize(s) => s.phf_hash(state),
61 ParsedKey::Bool(s) => s.phf_hash(state),
62 #[cfg(feature = "unicase")]
63 ParsedKey::UniCase(s) => s.phf_hash(state),
64 }
65 }
66}
67
68impl ParsedKey {
69 fn from_expr(expr: &Expr) -> Option<ParsedKey> {
70 match expr {
71 Expr::Lit(lit) => match &lit.lit {
72 Lit::Str(s) => Some(ParsedKey::Str(s.value())),
73 Lit::ByteStr(s) => Some(ParsedKey::Binary(s.value())),
74 Lit::Byte(s) => Some(ParsedKey::U8(s.value())),
75 Lit::Char(s) => Some(ParsedKey::Char(s.value())),
76 Lit::Int(s) => match s.suffix() {
77 "i8" => Some(ParsedKey::I8(s.base10_parse::<u8>().unwrap() as i8)),
81 "i16" => Some(ParsedKey::I16(s.base10_parse::<u16>().unwrap() as i16)),
82 "i32" => Some(ParsedKey::I32(s.base10_parse::<u32>().unwrap() as i32)),
83 "i64" => Some(ParsedKey::I64(s.base10_parse::<u64>().unwrap() as i64)),
84 "i128" => Some(ParsedKey::I128(s.base10_parse::<u128>().unwrap() as i128)),
85 "isize" => Some(ParsedKey::Isize(s.base10_parse::<usize>().unwrap() as isize)),
86 "u8" => Some(ParsedKey::U8(s.base10_parse::<u8>().unwrap())),
87 "u16" => Some(ParsedKey::U16(s.base10_parse::<u16>().unwrap())),
88 "u32" => Some(ParsedKey::U32(s.base10_parse::<u32>().unwrap())),
89 "u64" => Some(ParsedKey::U64(s.base10_parse::<u64>().unwrap())),
90 "u128" => Some(ParsedKey::U128(s.base10_parse::<u128>().unwrap())),
91 "usize" => Some(ParsedKey::Usize(s.base10_parse::<usize>().unwrap())),
92 _ => None,
93 },
94 Lit::Bool(s) => Some(ParsedKey::Bool(s.value)),
95 _ => None,
96 },
97 Expr::Array(array) => {
98 let mut buf = vec![];
99 for expr in &array.elems {
100 match expr {
101 Expr::Lit(lit) => match &lit.lit {
102 Lit::Int(s) => match s.suffix() {
103 "u8" | "" => buf.push(s.base10_parse::<u8>().unwrap()),
104 _ => return None,
105 },
106 _ => return None,
107 },
108 _ => return None,
109 }
110 }
111 Some(ParsedKey::Binary(buf))
112 }
113 Expr::Unary(unary) => {
114 macro_rules! try_negate (
118 ($val:expr) => {if $val < 0 { $val } else { -$val }}
119 );
120
121 match unary.op {
122 UnOp::Neg(_) => match ParsedKey::from_expr(&unary.expr)? {
123 ParsedKey::I8(v) => Some(ParsedKey::I8(try_negate!(v))),
124 ParsedKey::I16(v) => Some(ParsedKey::I16(try_negate!(v))),
125 ParsedKey::I32(v) => Some(ParsedKey::I32(try_negate!(v))),
126 ParsedKey::I64(v) => Some(ParsedKey::I64(try_negate!(v))),
127 ParsedKey::I128(v) => Some(ParsedKey::I128(try_negate!(v))),
128 ParsedKey::Isize(v) => Some(ParsedKey::Isize(try_negate!(v))),
129 _ => None,
130 },
131 UnOp::Deref(_) => {
132 let mut expr = &*unary.expr;
133 while let Expr::Group(group) = expr {
134 expr = &*group.expr;
135 }
136 match expr {
137 Expr::Lit(ExprLit {
138 lit: Lit::ByteStr(s),
139 ..
140 }) => Some(ParsedKey::Binary(s.value())),
141 _ => None,
142 }
143 }
144 _ => None,
145 }
146 }
147 Expr::Group(group) => ParsedKey::from_expr(&group.expr),
148 #[cfg(feature = "unicase")]
149 Expr::Call(call) => {
150 if let Expr::Path(ep) = call.func.as_ref() {
151 let segments = &mut ep.path.segments.iter().rev();
152 let last = &segments.next()?.ident;
153 let last_ahead = &segments.next()?.ident;
154 let is_unicode = last_ahead == "UniCase" && last == "unicode";
155 let is_ascii = last_ahead == "UniCase" && last == "ascii";
156 if call.args.len() == 1 && (is_unicode || is_ascii) {
157 if let Some(Expr::Lit(ExprLit {
158 attrs: _,
159 lit: Lit::Str(s),
160 })) = call.args.first()
161 {
162 let v = if is_unicode {
163 UniCase::unicode(s.value())
164 } else {
165 UniCase::ascii(s.value())
166 };
167 Some(ParsedKey::UniCase(v))
168 } else {
169 None
170 }
171 } else {
172 None
173 }
174 } else {
175 None
176 }
177 }
178 _ => None,
179 }
180 }
181}
182
183struct Key {
184 parsed: ParsedKey,
185 expr: Expr,
186}
187
188impl PhfHash for Key {
189 fn phf_hash<H>(&self, state: &mut H)
190 where
191 H: Hasher,
192 {
193 self.parsed.phf_hash(state)
194 }
195}
196
197impl Parse for Key {
198 fn parse(input: ParseStream<'_>) -> parse::Result<Key> {
199 let expr = input.parse()?;
200 let parsed = ParsedKey::from_expr(&expr)
201 .ok_or_else(|| Error::new_spanned(&expr, "unsupported key expression"))?;
202
203 Ok(Key { parsed, expr })
204 }
205}
206
207struct Entry {
208 key: Key,
209 value: Expr,
210}
211
212impl PhfHash for Entry {
213 fn phf_hash<H>(&self, state: &mut H)
214 where
215 H: Hasher,
216 {
217 self.key.phf_hash(state)
218 }
219}
220
221impl Parse for Entry {
222 fn parse(input: ParseStream<'_>) -> parse::Result<Entry> {
223 let key = input.parse()?;
224 input.parse::<Token![=>]>()?;
225 let value = input.parse()?;
226 Ok(Entry { key, value })
227 }
228}
229
230struct Map(Vec<Entry>);
231
232impl Parse for Map {
233 fn parse(input: ParseStream<'_>) -> parse::Result<Map> {
234 let parsed = Punctuated::<Entry, Token![,]>::parse_terminated(input)?;
235 let map = parsed.into_iter().collect::<Vec<_>>();
236 check_duplicates(&map)?;
237 Ok(Map(map))
238 }
239}
240
241struct Set(Vec<Entry>);
242
243impl Parse for Set {
244 fn parse(input: ParseStream<'_>) -> parse::Result<Set> {
245 let parsed = Punctuated::<Key, Token![,]>::parse_terminated(input)?;
246 let set = parsed
247 .into_iter()
248 .map(|key| Entry {
249 key,
250 value: syn::parse_str("()").unwrap(),
251 })
252 .collect::<Vec<_>>();
253 check_duplicates(&set)?;
254 Ok(Set(set))
255 }
256}
257
258fn check_duplicates(entries: &[Entry]) -> parse::Result<()> {
259 let mut keys = HashSet::new();
260 for entry in entries {
261 if !keys.insert(&entry.key.parsed) {
262 return Err(Error::new_spanned(&entry.key.expr, "duplicate key"));
263 }
264 }
265 Ok(())
266}
267
268fn build_map(entries: &[Entry], state: HashState) -> proc_macro2::TokenStream {
269 let key = state.key;
270 let disps = state.disps.iter().map(|&(d1, d2)| quote!((#d1, #d2)));
271 let entries = state.map.iter().map(|&idx| {
272 let key = &entries[idx].key.expr;
273 let value = &entries[idx].value;
274 quote!((#key, #value))
275 });
276
277 quote! {
278 phf::Map {
279 key: #key,
280 disps: &[#(#disps),*],
281 entries: &[#(#entries),*],
282 }
283 }
284}
285
286fn build_ordered_map(entries: &[Entry], state: HashState) -> proc_macro2::TokenStream {
287 let key = state.key;
288 let disps = state.disps.iter().map(|&(d1, d2)| quote!((#d1, #d2)));
289 let idxs = state.map.iter().map(|idx| quote!(#idx));
290 let entries = entries.iter().map(|entry| {
291 let key = &entry.key.expr;
292 let value = &entry.value;
293 quote!((#key, #value))
294 });
295
296 quote! {
297 phf::OrderedMap {
298 key: #key,
299 disps: &[#(#disps),*],
300 idxs: &[#(#idxs),*],
301 entries: &[#(#entries),*],
302 }
303 }
304}
305
306#[proc_macro]
307pub fn phf_map(input: TokenStream) -> TokenStream {
308 let map = parse_macro_input!(input as Map);
309 let state = phf_generator::generate_hash(&map.0);
310
311 build_map(&map.0, state).into()
312}
313
314#[proc_macro]
315pub fn phf_set(input: TokenStream) -> TokenStream {
316 let set = parse_macro_input!(input as Set);
317 let state = phf_generator::generate_hash(&set.0);
318
319 let map = build_map(&set.0, state);
320 quote!(phf::Set { map: #map }).into()
321}
322
323#[proc_macro]
324pub fn phf_ordered_map(input: TokenStream) -> TokenStream {
325 let map = parse_macro_input!(input as Map);
326 let state = phf_generator::generate_hash(&map.0);
327
328 build_ordered_map(&map.0, state).into()
329}
330
331#[proc_macro]
332pub fn phf_ordered_set(input: TokenStream) -> TokenStream {
333 let set = parse_macro_input!(input as Set);
334 let state = phf_generator::generate_hash(&set.0);
335
336 let map = build_ordered_map(&set.0, state);
337 quote!(phf::OrderedSet { map: #map }).into()
338}