phf_macros/
lib.rs

1//! A set of macros to generate Rust source for PHF data structures at compile time.
2//! See [the `phf` crate's documentation][phf] for details.
3//!
4//! [phf]: https://docs.rs/phf
5
6use 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                    // we've lost the sign at this point, so `-128i8` looks like `128i8`,
78                    // which doesn't fit in an `i8`; parse it as a `u8` and cast (to `0i8`),
79                    // which is handled below, by `Unary`
80                    "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                // if we received an integer literal (always unsigned) greater than i__::max_value()
115                // then casting it to a signed integer type of the same width will negate it to
116                // the same absolute value so we don't need to negate it here
117                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}