Skip to main content

leqx_macro/
lib.rs

1// SPDX-License-Identifier: MIT
2//! Macro for leqx
3use proc_macro::TokenStream;
4use quote::{quote, ToTokens};
5use syn::parse::{Parse, ParseStream};
6use syn::token::Brace;
7use syn::{Attribute, Expr, Receiver, Token, braced, parenthesized, parse_macro_input};
8
9mod kw {
10    syn::custom_keyword!(mode);
11    syn::custom_keyword!(sparse);
12    syn::custom_keyword!(dense);
13}
14
15#[derive(Copy, Clone)]
16enum DfaType {
17    Sparse,
18    Dense,
19}
20
21struct Pattern {
22    // meta: Vec<Attribute>,
23    // mode: Vec<syn::Ident>,
24    pattern_tok: syn::LitStr,
25    expr: syn::Expr,
26}
27
28struct Mode {
29    vis: syn::Visibility,
30    #[allow(dead_code)]
31    mode_kw: kw::mode,
32    name: syn::Ident,
33    embed: bool,
34    meta: Vec<Attribute>,
35    dfa_type: DfaType,
36    receiver: Receiver,
37    raw_ident: Option<syn::Ident>,
38    typ: syn::Type,
39    brace: Brace,
40    patterns: Vec<Pattern>,
41}
42
43struct Input {
44    context: syn::ItemStruct,
45    modes: Vec<Mode>,
46}
47
48impl Parse for DfaType {
49    fn parse(input: ParseStream) -> syn::Result<Self> {
50        if input.peek(kw::dense) {
51            input.parse::<kw::dense>().map(|_| DfaType::Dense)
52        } else if input.peek(kw::sparse) {
53            input.parse::<kw::sparse>().map(|_| DfaType::Sparse)
54        } else {
55            Err(input.error("Invalid DFA type; must be dense or sparse"))
56        }
57    }
58}
59
60impl Parse for Mode {
61    fn parse(input: ParseStream) -> syn::Result<Self> {
62        let meta = input.call(Attribute::parse_outer)?;
63        let mut saved_meta = vec![];
64        let mut embed = true;
65        let mut dfa_type = DfaType::Sparse;
66        for attr in meta {
67            if !attr.path().is_ident("leqxer") {
68                saved_meta.push(attr);
69                continue;
70            }
71            attr.parse_nested_meta(|meta| {
72                if meta.path.is_ident("embed") {
73                    embed = meta.value()?.parse::<syn::LitBool>()?.value;
74                    Ok(())
75                } else if meta.path.is_ident("dfa") {
76                    dfa_type = meta.value()?.parse()?;
77                    Ok(())
78                } else {
79                    Err(meta.error("Invalid leqxer option. Valid options are dfa={sparse,dense}, embed={true,false}"))
80                }
81            })?;
82        }
83        let vis = input.parse()?;
84        let mode_kw = input.parse()?;
85        let name = input.parse()?;
86
87        let (receiver, raw_ident) = if input.peek(syn::token::Paren) {
88            let content;
89            parenthesized!(content in input);
90
91            let receiver = content.parse()?;
92            content.parse::<Token![,]>()?;
93            (receiver, Some(content.parse()?))
94        } else {
95            (syn::parse_quote!(&self), None)
96        };
97
98        input.parse::<Token![->]>()?;
99
100        let typ = input.parse()?;
101
102        let content;
103        let brace = braced!(content in input);
104        let mut patterns = vec![];
105        while !content.is_empty() {
106            patterns.push(content.parse()?)
107        }
108
109        Ok(Mode {
110            meta: saved_meta,
111            vis,
112            mode_kw,
113            name,
114            receiver,
115            raw_ident,
116            typ,
117            brace,
118            patterns,
119            dfa_type,
120            embed,
121        })
122    }
123}
124
125impl Parse for Pattern {
126    fn parse(input: ParseStream) -> syn::Result<Self> {
127        let pattern_tok = <syn::LitStr as Parse>::parse(input)?;
128        <Token![=>]>::parse(input)?;
129        let expr = syn::Expr::parse(input)?;
130        let requires_comma = matches!(
131            expr,
132            Expr::If(_)
133                | Expr::Match(_)
134                | Expr::Block(_)
135                | Expr::Unsafe(_)
136                | Expr::While(_)
137                | Expr::Loop(_)
138                | Expr::ForLoop(_)
139                | Expr::TryBlock(_)
140                | Expr::Const(_)
141        );
142
143        if (requires_comma && !input.is_empty()) || input.peek(Token![,]) {
144            input.parse::<Token![,]>()?;
145        }
146        Ok(Pattern { pattern_tok, expr })
147    }
148}
149
150impl Parse for Input {
151    fn parse(input: ParseStream) -> syn::Result<Self> {
152        let mut result = Input {
153            context: input.parse()?,
154            modes: vec![],
155        };
156
157        while !input.is_empty() {
158            result.modes.push(input.parse()?);
159        }
160
161        Ok(result)
162    }
163}
164
165fn leqxer_impl(input: Input) -> Result<TokenStream, syn::Error> {
166    let mut mode_fns = proc_macro2::TokenStream::new();
167    for mode in input.modes.iter() {
168        let mode_span = mode.brace.span.join();
169        let input_sym = syn::Ident::new("input", proc_macro2::Span::mixed_site());
170        let mut pats = proc_macro2::TokenStream::new();
171        let mut pat_strs = vec![];
172        let mut matches = proc_macro2::TokenStream::new();
173        let extract = if let Some(ref ident) = mode.raw_ident {
174            quote! {
175                let #ident = std::str::from_utf8(&#input_sym.haystack()[start..end]).unwrap();
176            }
177        } else {
178            quote!()
179        };
180
181        for (i, pat) in mode.patterns.iter().enumerate() {
182            let re = &pat.pattern_tok;
183            let expr = &pat.expr;
184            pats.extend(quote! {#re,});
185            pat_strs.push(pat.pattern_tok.value());
186            matches.extend(quote! {#i => #expr,});
187        }
188
189        let (dfa, dfa_ty) = match mode.dfa_type {
190            DfaType::Sparse => {
191                let dfa = {
192                    use regex_automata::dfa::{StartKind, dense::DFA};
193                    DFA::builder()
194                        .configure(DFA::config().start_kind(StartKind::Anchored))
195                        .syntax(::regex_automata::util::syntax::Config::new()
196                            .dot_matches_new_line(true))
197                        .build_many(pat_strs.as_slice())
198                        .and_then(|dfa| dfa.to_sparse())
199                        .map_err(|err| syn::Error::new(mode.brace.span.join(), err))?
200                };
201                eprintln!("mode ({mode_name}) sparse size (le={le}. be={be})",
202                    mode_name = mode.name,
203                    le = dfa.to_bytes_little_endian().len(),
204                    be = dfa.to_bytes_big_endian().len(),
205                );
206                let construct = if mode.embed {
207                    let bstr_le =
208                        syn::LitByteStr::new(dfa.to_bytes_little_endian().as_slice(), mode_span);
209                    let bstr_be =
210                        syn::LitByteStr::new(dfa.to_bytes_big_endian().as_slice(), mode_span);
211                    quote! {
212                        unsafe{ ::regex_automata::dfa::sparse::DFA::from_bytes_unchecked(if cfg!(target_endian="big") { #bstr_be } else { #bstr_le }).unwrap().0 }
213                    }
214                } else {
215                    quote! {{
216                        use regex_automata::dfa::{dense::DFA, StartKind};
217                        DFA::builder()
218                            .configure(DFA::config().start_kind(StartKind::Anchored))
219                            .syntax(::regex_automata::util::syntax::Config::new()
220                                .dot_matches_new_line(true))
221                            .build_many(pat_strs.as_slice())
222                            .and_then(|dfa| dfa.to_sparse())
223                            .unwrap()
224                    }}
225                };
226                (construct, quote!{ ::regex_automata::dfa::sparse::DFA<&[u8]> })
227            }
228            DfaType::Dense => {
229                let dfa = {
230                    use regex_automata::dfa::{StartKind, dense::DFA};
231                    DFA::builder()
232                        .configure(DFA::config().start_kind(StartKind::Anchored))
233                        .syntax(::regex_automata::util::syntax::Config::new()
234                            .dot_matches_new_line(true))
235                        .build_many(pat_strs.as_slice())
236                        .map_err(|err| syn::Error::new(mode.brace.span.join(), err))
237                }?;
238                eprintln!("mode ({mode_name}) dense size (le={le}. be={be})",
239                          mode_name = mode.name,
240                          le = dfa.to_bytes_little_endian().0.len(),
241                          be = dfa.to_bytes_big_endian().0.len(),
242                );
243
244                let construct = if mode.embed {
245                    let bstr_le =
246                        syn::LitByteStr::new(dfa.to_bytes_little_endian().0.as_slice(), mode_span);
247                    let bstr_be =
248                        syn::LitByteStr::new(dfa.to_bytes_big_endian().0.as_slice(), mode_span);
249                    quote! {
250                        unsafe{ ::regex_automata::dfa::dense::DFA::from_bytes_unchecked(if cfg!(target_endian="big") { #bstr_be } else { #bstr_le }) }.unwrap().0
251                    }
252                } else {
253                    quote! {{
254                        use regex_automata::dfa::{dense::DFA, StartKind};
255                        DFA::builder()
256                            .configure(DFA::config().start_kind(StartKind::Anchored))
257                            .syntax(::regex_automata::util::syntax::Config::new()
258                                .dot_matches_new_line(true))
259                            .build_many(pat_strs.as_slice())
260                            .map_err(|err| syn::Error::new(mode.brace.span.join(), err))
261                            .unwrap()
262                    }}
263                };
264                (construct, quote!{ ::regex_automata::dfa::dense::DFA<&[u32]> })
265            }
266        };
267
268        let vis = &mode.vis;
269        let mode_name = &mode.name;
270        let mode_receiver = &mode.receiver;
271        let mode_ty = &mode.typ;
272        let mut mode_meta = proc_macro2::TokenStream::new();
273        for meta in &mode.meta {
274            meta.to_tokens(&mut mode_meta);
275        }
276        mode_fns.extend(quote! {
277            #mode_meta
278            #vis fn #mode_name(#mode_receiver, #input_sym: &mut ::leqx::__private::Input) -> Option<#mode_ty> {
279                use ::regex_automata::dfa::Automaton;
280                static RE_SET: ::leqx::__private::LazyLock<#dfa_ty> = ::leqx::__private::LazyLock::new(|| #dfa);
281                #input_sym.set_anchored(::regex_automata::Anchored::Yes);
282                let res = RE_SET.try_search_fwd(#input_sym).unwrap()?;
283                let start = #input_sym.start();
284                let end = res.offset();
285                #input_sym.set_start(end);
286                #extract
287                Some(match res.pattern().as_usize() {
288                    #matches
289                    _ => unreachable!(),
290                })
291            }
292        });
293    }
294
295    let struct_decl = &input.context;
296    let (struct_impl, struct_type_params, struct_where) = input.context.generics.split_for_impl();
297    let struct_name = &input.context.ident;
298    Ok(quote! {
299        #struct_decl
300        impl #struct_impl #struct_name #struct_type_params #struct_where {
301            #mode_fns
302        }
303    }
304    .into())
305}
306
307#[proc_macro]
308pub fn leqxer(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
309    let input = parse_macro_input!(input as Input);
310    let result = leqxer_impl(input).unwrap_or_else(|err| err.into_compile_error().into());
311
312    result
313}