const_format_proc_macros/
parse_utils.rs

1use crate::{spanned::Spans, utils::Peekable2, Error};
2
3use proc_macro2::{
4    token_stream::IntoIter, Delimiter, Group, Ident, Punct, Span, TokenStream as TokenStream2,
5    TokenTree as TokenTree2,
6};
7
8use std::{cmp::PartialEq, ops::Range};
9
10pub type ParseStream<'a> = &'a mut ParseBuffer;
11
12pub struct ParseBuffer {
13    iter: Peekable2<IntoIter>,
14}
15
16impl ParseBuffer {
17    pub fn new(ts: TokenStream2) -> Self {
18        let iter = Peekable2::new(ts);
19        Self { iter }
20    }
21
22    pub fn is_empty(&mut self) -> bool {
23        self.iter.is_empty()
24    }
25
26    pub fn peek(&mut self) -> Option<&TokenTree2> {
27        self.iter.peek()
28    }
29    pub fn peek2(&mut self) -> Option<&TokenTree2> {
30        self.iter.peek2()
31    }
32
33    pub fn parse_punct(&mut self, c: char) -> Result<Punct, crate::Error> {
34        match self.next() {
35            Some(TokenTree2::Punct(x)) if x.as_char() == c => Ok(x),
36            Some(x) => Err(Error::new(x.span(), &format!("Expected a '{}' token", c))),
37            None => Err(Error::new(
38                Span::mixed_site(),
39                &format!("Expected a '{}' token", c),
40            )),
41        }
42    }
43    pub fn parse_opt_punct(&mut self, c: char) -> Result<Option<Punct>, crate::Error> {
44        match self.next() {
45            Some(TokenTree2::Punct(x)) if x.as_char() == c => Ok(Some(x)),
46            Some(x) => Err(Error::new(x.span(), &format!("Expected a '{}' token", c))),
47            None => Ok(None),
48        }
49    }
50
51    pub fn parse_ident(&mut self) -> Result<Ident, crate::Error> {
52        match self.next() {
53            Some(TokenTree2::Ident(x)) => Ok(x),
54            Some(x) => Err(Error::new(x.span(), "Expected an identifier")),
55            None => Err(Error::new(Span::mixed_site(), "Expected an identifier")),
56        }
57    }
58
59    pub fn parse_paren(&mut self) -> Result<Parentheses, crate::Error> {
60        match self.next() {
61            Some(TokenTree2::Group(group)) if group.delimiter() == Delimiter::Parenthesis => {
62                Ok(Parentheses {
63                    paren_span: group.span(),
64                    contents: group.stream(),
65                })
66            }
67            Some(x) => Err(Error::new(
68                x.span(),
69                &format!("Expected parentheses: found {}", x),
70            )),
71            None => Err(Error::new(
72                Span::mixed_site(),
73                "Expected parentheses, found nothing",
74            )),
75        }
76    }
77
78    pub fn parse_unwrap_paren<F, T>(&mut self, f: F) -> Result<T, crate::Error>
79    where
80        F: FnOnce(ParseStream<'_>) -> Result<T, crate::Error>,
81    {
82        if matches!(self.peek(), Some(TokenTree2::Group(x)) if x.delimiter() == Delimiter::Parenthesis )
83        {
84            if let Some(TokenTree2::Group(group)) = self.next() {
85                ParseBuffer::new(group.stream()).parse_unwrap_tt(f)
86            } else {
87                unreachable!("But I peeked for a Parenthesis delimited TokenTree::Group!!")
88            }
89        } else {
90            f(self)
91        }
92    }
93
94    pub fn parse_unwrap_group<F, T>(&mut self, f: F) -> Result<T, crate::Error>
95    where
96        F: FnOnce(ParseStream<'_>) -> Result<T, crate::Error>,
97    {
98        if let Some(TokenTree2::Group(group)) = self.next() {
99            ParseBuffer::new(group.stream()).parse_unwrap_tt(f)
100        } else {
101            f(self)
102        }
103    }
104
105    pub fn parse_token_stream_and_span(&mut self) -> (TokenStream2, Spans) {
106        let mut start = match self.peek() {
107            Some(x) => x.span(),
108            None => Span::call_site(),
109        };
110
111        let mut end = start;
112
113        let ts = self
114            .inspect(|tt| {
115                end = tt.span();
116                if let Some(next) = start.join(end) {
117                    start = next;
118                }
119            })
120            .collect::<TokenStream2>();
121
122        (ts, Spans { start, end })
123    }
124
125    /// Unwraps a none-delimited token tree to parse a type,
126    /// if the first token is not a none-delimited token tree it parses the type in
127    /// the passed in ParseStream.
128    pub fn parse_unwrap_tt<F, T>(&mut self, f: F) -> Result<T, crate::Error>
129    where
130        F: FnOnce(ParseStream<'_>) -> Result<T, crate::Error>,
131    {
132        if matches!(self.peek(), Some(TokenTree2::Group(x)) if x.delimiter() == Delimiter::None ) {
133            if let Some(TokenTree2::Group(group)) = self.next() {
134                ParseBuffer::new(group.stream()).parse_unwrap_tt(f)
135            } else {
136                unreachable!("But I peeked for a None delimited TokenTree::Group!!")
137            }
138        } else {
139            f(self)
140        }
141    }
142}
143
144impl Iterator for ParseBuffer {
145    type Item = TokenTree2;
146
147    fn next(&mut self) -> Option<TokenTree2> {
148        self.iter.next()
149    }
150
151    fn size_hint(&self) -> (usize, Option<usize>) {
152        self.iter.size_hint()
153    }
154}
155
156///////////////////////////////////////////////////////////////////////////////
157
158pub struct Parentheses {
159    #[allow(dead_code)]
160    pub paren_span: Span,
161    pub contents: TokenStream2,
162}
163
164///////////////////////////////////////////////////////////////////////////////
165
166pub struct LitStr {
167    value: String,
168    pub rawness: StrRawness,
169    pub inside_lit: Range<usize>,
170    pub span: Span,
171}
172
173impl LitStr {
174    pub fn value(&self) -> &str {
175        &self.value[self.inside_lit.clone()]
176    }
177    pub(crate) fn parse_from_literal(literal: &proc_macro2::Literal) -> Result<Self, Error> {
178        let mut value = literal.to_string();
179        // Ignoring the quote characters
180        let mut range = 1..value.len() - 1;
181        let span = literal.span();
182
183        let is_raw = if let Some(suffix) = value.strip_prefix('r') {
184            let hashes = suffix.bytes().take_while(|x| *x == b'#').count();
185
186            if value.as_bytes()[1 + hashes] != b'"' {
187                return Err(Error::new(
188                    span,
189                    &format!("Expected a string literal, found: {}", literal),
190                ));
191            }
192
193            // Ignoring the r and hashes
194            range.start += 1 + hashes;
195            range.end -= hashes;
196            Some(hashes as u32)
197        } else {
198            let mut matches = value.match_indices(r#"\u"#).peekable();
199            if matches.peek().is_some() {
200                let mut prev_end = 0;
201                let mut new = String::with_capacity(value.len());
202
203                for (pos, _) in matches {
204                    new.push_str(&value[prev_end..pos]);
205
206                    let past_open = pos + 3;
207
208                    let off_close = value[pos..].find('}').unwrap();
209
210                    let c = &value[past_open..pos + off_close];
211                    let c = u32::from_str_radix(c, 16).unwrap();
212                    let c = std::char::from_u32(c).unwrap();
213
214                    // if matches!(c, '\\' | '"') {
215                    //     new.push('\\');
216                    // }
217                    new.push(c);
218
219                    prev_end = pos + off_close + 1;
220                }
221                new.push_str(&value[prev_end..]);
222                value = new;
223            }
224
225            range = 1..value.len() - 1;
226
227            None
228        };
229
230        Ok(Self {
231            value,
232            rawness: StrRawness { is_raw, span },
233            inside_lit: range,
234            span,
235        })
236    }
237}
238
239#[derive(Debug, Copy, Clone)]
240pub struct StrRawness {
241    is_raw: Option<u32>,
242    span: Span,
243}
244
245impl PartialEq for StrRawness {
246    fn eq(&self, other: &Self) -> bool {
247        self.is_raw == other.is_raw
248    }
249}
250
251impl StrRawness {
252    #[cfg(test)]
253    pub fn dummy() -> Self {
254        Self {
255            is_raw: Some(4),
256            span: Span::mixed_site(),
257        }
258    }
259
260    pub fn span(&self) -> Span {
261        self.span
262    }
263
264    /// Tokenizes a slice of the parsed string literal.
265    pub fn tokenize_sub(&self, str: &str) -> TokenStream2 {
266        let mut buffer = String::new();
267        match self.is_raw {
268            Some(hashes) => {
269                let hashes = hashes as usize;
270                buffer.reserve(3 + hashes + str.len() + hashes);
271                buffer.push('r');
272                let hashes = (0..hashes).map(|_| '#');
273                buffer.extend(hashes.clone());
274                buffer.push('"');
275                buffer.push_str(str);
276                buffer.push('"');
277                buffer.extend(hashes);
278            }
279            None => {
280                buffer.reserve(2 + str.len());
281                buffer.push('"');
282                buffer.push_str(str);
283                buffer.push('"');
284            }
285        }
286
287        buffer
288            .parse::<TokenStream2>()
289            .unwrap()
290            .set_span_recursive(self.span)
291    }
292}
293
294///////////////////////////////////////////////////////////////////////////////
295
296pub trait TokenTreeExt: Sized {
297    fn as_token_tree(&self) -> &TokenTree2;
298    fn into_token_tree(self) -> TokenTree2;
299
300    fn is_punct(&self, c: char) -> bool {
301        matches!(self.as_token_tree(), TokenTree2::Punct(p)  if p.as_char() == c)
302    }
303
304    #[allow(dead_code)]
305    fn is_paren(&self) -> bool {
306        matches!(
307            self.as_token_tree(),
308            TokenTree2::Group(g) if g.delimiter() == Delimiter::Parenthesis
309        )
310    }
311
312    #[allow(dead_code)]
313    fn is_ident(&self, ident: &str) -> bool {
314        matches!(self.as_token_tree(), TokenTree2::Ident(x)  if x == ident)
315    }
316
317    fn set_span_recursive(self, span: Span) -> TokenTree2 {
318        let mut tt = self.into_token_tree();
319
320        tt.set_span(span);
321        if let TokenTree2::Group(group) = tt {
322            let delim = group.delimiter();
323            let stream = group.stream().set_span_recursive(span);
324            tt = TokenTree2::Group(Group::new(delim, stream));
325        }
326        tt.set_span(span);
327        tt
328    }
329}
330
331impl TokenTreeExt for TokenTree2 {
332    fn as_token_tree(&self) -> &TokenTree2 {
333        self
334    }
335
336    fn into_token_tree(self) -> TokenTree2 {
337        self
338    }
339}
340
341///////////////////////////////////////////////////////////////////////////////
342
343pub trait TokenStream2Ext: Sized {
344    fn into_token_stream(self) -> TokenStream2;
345
346    fn set_span_recursive(self, span: Span) -> TokenStream2 {
347        self.into_token_stream()
348            .into_iter()
349            .map(|tt| tt.set_span_recursive(span))
350            .collect()
351    }
352}
353
354impl TokenStream2Ext for TokenStream2 {
355    fn into_token_stream(self) -> TokenStream2 {
356        self
357    }
358}
359
360///////////////////////////////////////////////////////////////////////////////
361
362pub trait MyParse: Sized {
363    fn parse(input: ParseStream<'_>) -> Result<Self, crate::Error>;
364
365    fn parse_token_stream_1(input: proc_macro::TokenStream) -> Result<Self, crate::Error> {
366        Self::parse(&mut ParseBuffer::new(TokenStream2::from(input)))
367    }
368
369    #[allow(dead_code)]
370    fn parse_token_stream_2(input: TokenStream2) -> Result<Self, crate::Error> {
371        Self::parse(&mut ParseBuffer::new(input))
372    }
373}
374
375///////////////////////////////////////////////////////////////////////////////