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

proc_macro_tool/
parse_iter.rs
