unfmt_macros/
unfmt_macros.rs

1#![no_std]
2
3extern crate alloc;
4extern crate proc_macro;
5
6use alloc::{borrow::ToOwned as _, format, vec::Vec};
7
8use bstr::ByteSlice as _;
9use proc_macro::TokenStream;
10use proc_macro2::Span;
11use quote::{quote, ToTokens};
12use syn::{
13    parse::{Parse, ParseStream, Result},
14    parse_macro_input, parse_str, Expr, Ident, Lit, LitBool, LitByteStr, Token, TypePath,
15};
16
17struct Unformat {
18    full_match: bool,
19    is_pattern_str: bool,
20    pattern: Vec<u8>,
21    text: Expr,
22}
23
24impl Parse for Unformat {
25    fn parse(input: ParseStream) -> Result<Self> {
26        #[expect(
27            clippy::wildcard_enum_match_arm,
28            reason = "We want to match on future variants as well."
29        )]
30        let (pattern, is_pattern_str) = match input.parse::<Lit>()? {
31            Lit::Str(str) => (str.value().into_bytes(), true),
32            Lit::ByteStr(byte_str) => (byte_str.value(), false),
33            _ => return Err(input.error("expected a string literal")),
34        };
35
36        input.parse::<Token![,]>()?;
37
38        let text = input.parse::<Expr>()?;
39
40        let full_match = if input.parse::<Token![,]>().is_ok() {
41            input.parse::<LitBool>().is_ok_and(|bool| bool.value)
42        } else {
43            false
44        };
45        Ok(Self {
46            full_match,
47            is_pattern_str,
48            pattern,
49            text,
50        })
51    }
52}
53
54enum Assignee {
55    Index(u32),
56    Variable(Ident),
57}
58
59impl Assignee {
60    fn new(variable: &str, index: &mut u32) -> Self {
61        variable.parse::<u32>().map_or_else(
62            |_| {
63                if variable.is_empty() {
64                    let tuple_index = *index;
65                    *index = index.saturating_add(1);
66                    Self::Index(tuple_index)
67                } else {
68                    Self::Variable(parse_str(variable).expect("invalid variable name"))
69                }
70            },
71            Self::Index,
72        )
73    }
74}
75
76enum CaptureTypePath {
77    Bytes,
78    Str,
79    Typed(TypePath),
80}
81
82impl CaptureTypePath {
83    fn new(type_path: &str, is_pattern_str: bool) -> Self {
84        if type_path.is_empty() {
85            if is_pattern_str {
86                Self::Str
87            } else {
88                Self::Bytes
89            }
90        } else if type_path == "&str" {
91            Self::Str
92        } else if type_path == "&[u8]" {
93            Self::Bytes
94        } else {
95            Self::Typed(parse_str(type_path).expect("invalid type path"))
96        }
97    }
98}
99
100impl ToTokens for CaptureTypePath {
101    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
102        tokens.extend(match *self {
103            Self::Str => {
104                quote! { &str }
105            }
106            Self::Bytes => {
107                quote! { &[u8] }
108            }
109            Self::Typed(ref type_path) => {
110                quote! { #type_path }
111            }
112        });
113    }
114}
115
116struct Capture {
117    assignee: Assignee,
118    text: Vec<u8>,
119    r#type: CaptureTypePath,
120}
121
122impl Capture {
123    fn new(text: &[u8], capture: &str, is_pattern_str: bool, index: &mut u32) -> Self {
124        let (variable, type_path) = capture.split_once(':').unwrap_or((capture, ""));
125        Self {
126            text: text.to_vec(),
127            assignee: Assignee::new(variable, index),
128            r#type: CaptureTypePath::new(type_path, is_pattern_str),
129        }
130    }
131}
132
133impl ToTokens for Capture {
134    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
135        let rhs = match self.r#type {
136            CaptureTypePath::Str => {
137                quote! {
138                    if let Ok(__unfmt_left) = __unfmt_left.to_str() {
139                        __unfmt_left
140                    } else {
141                        break 'unformat None;
142                    }
143                }
144            }
145            CaptureTypePath::Bytes => {
146                quote! { __unfmt_left }
147            }
148            CaptureTypePath::Typed(ref type_path) => {
149                quote! {
150                    if let Ok(Ok(__unfmt_left)) = __unfmt_left.to_str().map(|value| value.parse::<#type_path>()) {
151                        __unfmt_left
152                    } else {
153                        break 'unformat None;
154                    }
155                }
156            }
157        };
158        let assignment = match self.assignee {
159            Assignee::Index(ref index) => {
160                let ident = Ident::new(&format!("__unfmt_capture_{index}"), Span::call_site());
161                quote! { let #ident = #rhs }
162            }
163            Assignee::Variable(ref ident) => {
164                quote! { #ident = Some(#rhs) }
165            }
166        };
167        let text = LitByteStr::new(&self.text, Span::call_site());
168
169        // If text is empty, `find` will return `Some(0)` and the capture will
170        // be at the end of the pattern, so this capture (`__unfmt_left`) would
171        // be empty. Since captures are inherently .*? in regex, this capture
172        // should consume the remainder of the text, so we swap `__unfmt_left`
173        // and `__unfmt_right` to achieve this.
174        tokens.extend(if self.text.is_empty() {
175            quote! { let (__unfmt_left, __unfmt_right) = (__unfmt_byte_text, b""); }
176        } else {
177            quote! {
178                let Some((__unfmt_left, __unfmt_right)) = __unfmt_byte_text.split_once_str(#text) else {
179                    break 'unformat None;
180                };
181            }
182        });
183
184        tokens.extend(quote! {
185            #assignment;
186            __unfmt_byte_text = BStr::new(__unfmt_right);
187        });
188    }
189}
190
191/// Basic implementation of reversing the `format!` process. Matches a given
192/// text against a given pattern, returning any captures.
193///
194/// Rules:
195///
196///  - Patterns are substring matched.
197///  - Captures are written as `{<index-or-var>?(:<type>)?}` in the pattern.
198///  - Captures are similar to `(.*?)` in regex, but without backtracking.
199///  - Sequential captures (e.g. `{}{}`) are not supported and will return
200///    `None`.
201///
202/// # Panics
203///
204/// This function panics if the pattern is invalid. This includes:
205///
206///  - Consecutive captures.
207///  - Unmatched `}` in the pattern.
208///  - Invalid UTF-8 in capture names.
209///
210#[proc_macro]
211#[inline]
212pub fn unformat(input: TokenStream) -> TokenStream {
213    let Unformat {
214        pattern,
215        text,
216        is_pattern_str,
217        full_match,
218    } = parse_macro_input!(input as Unformat);
219
220    let (initial_part, captures) = compile(&pattern, is_pattern_str);
221    let initial_part = Lit::ByteStr(LitByteStr::new(&initial_part, Span::call_site()));
222
223    let capture_idents = {
224        let mut capture_indices = captures
225            .iter()
226            .filter_map(|capture| match capture.assignee {
227                Assignee::Index(capture_index) => Some(capture_index),
228                Assignee::Variable(..) => None,
229            })
230            .collect::<Vec<_>>();
231
232        capture_indices.sort_by_key(|&index| index);
233
234        capture_indices
235            .into_iter()
236            .map(|index| Ident::new(&format!("__unfmt_capture_{index}"), Span::call_site()))
237            .collect::<Vec<_>>()
238    };
239
240    let capture_block = if full_match {
241        quote! {
242            if !__unfmt_left.is_empty() {
243                break 'unformat None;
244            }
245            #(#captures)*
246            if !__unfmt_byte_text.is_empty() {
247                break 'unformat None;
248            }
249        }
250    } else {
251        quote! { #(#captures)* }
252    };
253
254    TokenStream::from(quote! {
255        'unformat: {
256            use ::core::str::FromStr;
257            use ::unfmt::bstr::{ByteSlice, BStr};
258            let Some((__unfmt_left, mut __unfmt_byte_text)) = BStr::new(#text).split_once_str(#initial_part) else {
259                break 'unformat None;
260            };
261            #capture_block
262            Some((#(#capture_idents),*))
263        }
264    })
265}
266
267fn compile(pattern: &[u8], is_pattern_str: bool) -> (Vec<u8>, Vec<Capture>) {
268    let mut pattern = pattern.replace(b"{{", "\u{f8fd}");
269    pattern.reverse();
270    let mut pattern = pattern.replace(b"}}", "\u{f8fe}");
271    pattern.reverse();
272
273    let mut pattern_parts = pattern.split_str("{");
274
275    // SAFETY: The first part is always present.
276    let initial_part = unsafe {
277        pattern_parts
278            .next()
279            .unwrap_unchecked()
280            .replace("\u{f8fd}", "{")
281    };
282
283    let mut current_index: u32 = 0;
284    let mut compiled_pattern = Vec::new();
285    for pattern_part in pattern_parts {
286        let (capture, text) = pattern_part
287            .split_once_str("}")
288            .expect("unmatched } in pattern");
289        let capture = capture
290            .to_str()
291            .expect("invalid UTF-8 in capture names")
292            .to_owned();
293        let mut text = text.replace("\u{f8fd}", b"{");
294        text.reverse();
295        let mut text = text.replace("\u{f8fe}", b"}");
296        text.reverse();
297        compiled_pattern.push(Capture::new(
298            &text,
299            &capture,
300            is_pattern_str,
301            &mut current_index,
302        ));
303    }
304
305    assert!(
306        compiled_pattern.windows(2).all(|parts| parts
307            .iter()
308            .any(|&Capture { ref text, .. }| !text.is_empty())),
309        "consecutive captures are not allowed"
310    );
311
312    (initial_part, compiled_pattern)
313}
unfmt_macros/unfmt_macros.rs

unfmt_macros/
unfmt_macros.rs
