proc_macro_tool/
parse_iter.rs

1use crate::{stream, TokenTreeExt as _};
2use proc_macro::{Group, Spacing::*, Span, TokenStream, TokenTree};
3use std::{array, collections::VecDeque, iter::{self, FusedIterator}};
4
5/// Create [`ParseIter`]
6pub trait ParseIterExt: IntoIterator<Item = TokenTree> + Sized {
7    /// Create [`ParseIter`]
8    fn parse_iter(self) -> ParseIter<Self::IntoIter> {
9        ParseIter {
10            iter: self.into_iter(),
11            buf: VecDeque::new(),
12            end_span: Span::call_site(),
13        }
14    }
15}
16impl<I: IntoIterator<Item = TokenTree>> ParseIterExt for I { }
17
18/// Create [`ParseIter`]
19pub trait ParseIterGroupExt: Into<Group> + Sized {
20    /// Create [`ParseIter`]
21    fn stream_parse_iter(self) -> ParseIter<<TokenStream as IntoIterator>::IntoIter> {
22        let group: Group = self.into();
23        ParseIter {
24            iter: group.stream().into_iter(),
25            buf: VecDeque::new(),
26            end_span: group.span_close(),
27        }
28    }
29}
30impl ParseIterGroupExt for Group { }
31
32/// Peek `n` iterator adapter
33#[derive(Debug, Clone)]
34pub struct ParseIter<I: Iterator<Item = TokenTree>> {
35    iter: I,
36    buf: VecDeque<TokenTree>,
37    end_span: Span,
38}
39
40impl<I: Iterator<Item = TokenTree>> ParseIter<I> {
41    pub fn peek(&mut self) -> Option<&TokenTree> {
42        self.peek_i(0)
43    }
44
45    pub fn next_if<F>(&mut self, f: F) -> Option<TokenTree>
46    where F: FnOnce(&TokenTree) -> bool,
47    {
48        let peek = self.peek()?;
49
50        if f(peek) {
51            self.next()
52        } else {
53            None
54        }
55    }
56
57    pub fn next_is<F>(&mut self, f: F) -> bool
58    where F: FnOnce(TokenTree) -> bool,
59    {
60        self.next().is_some_and(f)
61    }
62
63    pub fn next_i_if<F>(&mut self, i: usize, f: F) -> Option<TokenTree>
64    where F: FnOnce(&TokenTree) -> bool,
65    {
66        let peek = self.peek_i(i)?;
67
68        if f(peek) {
69            self.nth(i)
70        } else {
71            None
72        }
73    }
74
75    /// # Panics
76    /// `self.count() < N`
77    #[track_caller]
78    pub fn next_tts<const N: usize>(&mut self) -> [TokenTree; N] {
79        array::from_fn(|_| self.next()
80            .expect("unexpected end of input"))
81    }
82
83    pub fn peek_i(&mut self, i: usize) -> Option<&TokenTree> {
84        for _ in self.buf.len()..=i {
85            self.buf.push_back(self.iter.next()?);
86        }
87        Some(&self.buf[i])
88    }
89
90    /// Get current token span.
91    /// return [`self.end_span`](#method.end_span) when nothing token
92    pub fn span(&mut self) -> Span {
93        let end_span = self.end_span;
94        self.peek().map_or(end_span, TokenTree::span)
95    }
96
97    /// Get `i`th token span.
98    /// return [`self.end_span`](#method.end_span) when nothing token
99    pub fn span_i(&mut self, i: usize) -> Span {
100        let end_span = self.end_span;
101        self.peek_i(i).map_or(end_span, TokenTree::span)
102    }
103
104    pub fn peek_is<F>(&mut self, f: F) -> bool
105    where F: FnOnce(&TokenTree) -> bool,
106    {
107        self.peek().is_some_and(f)
108    }
109
110    pub fn peek_i_is<F>(&mut self, i: usize, f: F) -> bool
111    where F: FnOnce(&TokenTree) -> bool,
112    {
113        self.peek_i(i).is_some_and(f)
114    }
115
116    /// Peek jointed puncts
117    pub fn peek_puncts(
118        &mut self,
119        puncts: impl AsRef<[u8]>,
120    ) -> Option<impl Iterator<Item = &TokenTree>> {
121        self.peek_i_puncts(0, puncts)
122    }
123
124    /// Peek jointed puncts
125    pub fn peek_i_puncts(
126        &mut self,
127        i: usize,
128        puncts: impl AsRef<[u8]>,
129    ) -> Option<impl Iterator<Item = &TokenTree>> {
130        let mut prev = None;
131
132        let puncts = puncts.as_ref();
133
134        for (j, ch) in puncts.iter().copied().map(char::from).enumerate() {
135            if let Some(prev) = prev {
136                if prev == Alone { return None }
137            }
138
139            let tt = self.peek_i(i + j)?;
140            let TokenTree::Punct(p) = tt else { return None };
141            if p.as_char() != ch { return None }
142
143            prev = Some(p.spacing());
144        }
145        Some(self.buf.iter().skip(i).take(puncts.len()))
146    }
147
148    pub fn is_puncts(&mut self, puncts: impl AsRef<[u8]>) -> bool {
149        self.peek_puncts(puncts).is_some()
150    }
151
152    pub fn is_puncts_at(&mut self, i: usize, puncts: impl AsRef<[u8]>) -> bool {
153        self.peek_i_puncts(i, puncts).is_some()
154    }
155
156    /// Next jointed puncts
157    pub fn next_puncts(
158        &mut self,
159        puncts: impl AsRef<[u8]>,
160    ) -> Option<impl Iterator<Item = TokenTree> + '_> {
161        let puncts = puncts.as_ref();
162        let _ = self.peek_puncts(puncts)?;
163
164        Some(self.buf.drain(..puncts.len()))
165    }
166
167    /// Split [`TokenStream`] with `puncts`
168    ///
169    /// Like `"+-,-+".split_puncts(",")` -> `"+-"`
170    ///
171    /// `self` is `"-+"`
172    #[allow(clippy::missing_panics_doc)]
173    pub fn split_puncts(&mut self, puncts: impl AsRef<[u8]>) -> Option<TokenStream> {
174        let puncts = puncts.as_ref();
175        let mut i = 0;
176
177        loop {
178            if self.peek_i_puncts(i, puncts).is_some() {
179                let left = self.take(i).collect();
180                let _ = self.next_puncts(puncts).unwrap();
181                break Some(left);
182            }
183            self.peek_i(i)?;
184            i += 1;
185        }
186    }
187
188    /// Split [`TokenStream`] with `puncts`
189    ///
190    /// Like `"+-,-+".split_puncts_exclude(",")` -> `"+-"`
191    ///
192    /// `self` is `",-+"`
193    #[allow(clippy::missing_panics_doc)]
194    pub fn split_puncts_exclude(&mut self, puncts: impl AsRef<[u8]>) -> Option<TokenStream> {
195        let puncts = puncts.as_ref();
196        let mut i = 0;
197
198        loop {
199            if self.peek_i_puncts(i, puncts).is_some() {
200                let left = self.take(i).collect();
201                break Some(left);
202            }
203            self.peek_i(i)?;
204            i += 1;
205        }
206    }
207
208    /// Split [`TokenStream`] with `puncts`
209    ///
210    /// Like `"+-,-+".split_puncts_include(",")` -> `"+-,"`
211    ///
212    /// `self` is `"-+"`
213    #[allow(clippy::missing_panics_doc)]
214    pub fn split_puncts_include(&mut self, puncts: impl AsRef<[u8]>) -> Option<TokenStream> {
215        let puncts = puncts.as_ref();
216        let mut i = 0;
217
218        loop {
219            if self.peek_i_puncts(i, puncts).is_some() {
220                let left = self.take(i+puncts.len()).collect();
221                break Some(left);
222            }
223            self.peek_i(i)?;
224            i += 1;
225        }
226    }
227
228    /// Split all [`TokenStream`] with `puncts`
229    ///
230    /// Like `"+-,-+".split_puncts(",")` -> `"+-", "-+"`
231    pub fn split_puncts_all<'a>(
232        &'a mut self,
233        puncts: impl AsRef<[u8]> + 'a,
234    ) -> impl Iterator<Item = TokenStream> + 'a {
235        iter::from_fn(move || {
236            self.split_puncts(puncts.as_ref())
237                .or_else(|| self.peek().is_some().then(|| self.collect()))
238        })
239    }
240
241    /// # Note
242    ///
243    /// This method parse like `#[...]`
244    pub fn next_attributes(&mut self) -> Vec<TokenStream> {
245        let mut attributes = vec![];
246
247        while self.peek_puncts("#").is_some()
248            &&self.peek_i_is(1, |tt| tt.is_delimiter_bracket())
249        {
250            attributes.push(self.next_tts::<2>().into_iter().collect());
251        }
252
253        attributes
254    }
255
256    pub fn push_if_to<To, F>(&mut self, to: &mut To, f: F) -> bool
257    where To: Extend<TokenTree>,
258          F: FnOnce(&TokenTree) -> bool,
259    {
260        self.peek()
261            .is_some_and(f)
262            .then(|| to.extend(self.next()))
263            .is_some()
264    }
265
266    /// # Note
267    ///
268    /// This method parse like `#![...]`
269    pub fn next_outer_attributes(&mut self) -> Vec<TokenStream> {
270        let mut attributes = vec![];
271
272        while self.peek_puncts("#").is_some()
273            &&self.peek_i_is(1, |tt| tt.is_punch('!'))
274            &&self.peek_i_is(2, |tt| tt.is_delimiter_bracket())
275        {
276            attributes.push(stream(self.next_tts::<3>()));
277        }
278
279        attributes
280    }
281
282    #[allow(clippy::missing_panics_doc)]
283    pub fn next_vis(&mut self) -> Option<TokenStream> {
284        if self.peek_is(|tt| tt.is_keyword("pub")) {
285            if self.peek_i_is(1, |tt| tt.is_delimiter_paren()) {
286                return Some(stream(self.next_tts::<2>()));
287            }
288            return Some(self.next().unwrap().into());
289        }
290        None
291    }
292
293    pub fn end_span(&self) -> Span {
294        self.end_span
295    }
296
297    pub fn set_end_span(&mut self, end_span: Span) {
298        self.end_span = end_span;
299    }
300}
301impl<I: Iterator<Item = TokenTree>> Iterator for ParseIter<I> {
302    type Item = TokenTree;
303
304    fn next(&mut self) -> Option<Self::Item> {
305        self.buf.pop_front()
306            .or_else(|| self.iter.next())
307    }
308
309    fn count(self) -> usize
310    where Self: Sized,
311    {
312        self.buf.len() + self.iter.count()
313    }
314
315    fn fold<B, F>(self, init: B, f: F) -> B
316    where Self: Sized,
317          F: FnMut(B, Self::Item) -> B,
318    {
319        self.buf.into_iter().chain(self.iter).fold(init, f)
320    }
321
322    fn size_hint(&self) -> (usize, Option<usize>) {
323        let (lo, hi) = self.iter.size_hint();
324        let lo = lo.saturating_add(self.buf.len());
325        let hi = hi.and_then(|hi| hi.checked_add(self.buf.len()));
326        (lo, hi)
327    }
328}
329impl<I: DoubleEndedIterator<Item = TokenTree>> DoubleEndedIterator for ParseIter<I> {
330    fn next_back(&mut self) -> Option<Self::Item> {
331        self.iter.next_back()
332            .or_else(|| self.buf.pop_back())
333    }
334
335    fn rfold<B, F>(self, init: B, f: F) -> B
336    where Self: Sized,
337          F: FnMut(B, Self::Item) -> B,
338    {
339        self.buf.into_iter().chain(self.iter).rfold(init, f)
340    }
341}
342impl<I: ExactSizeIterator<Item = TokenTree>> ExactSizeIterator for ParseIter<I> { }
343impl<I: FusedIterator<Item = TokenTree>> FusedIterator for ParseIter<I> { }
proc_macro_tool/parse_iter.rs

proc_macro_tool/
parse_iter.rs
