proc_macro_tool/
parse_iter.rs

1use crate::{stream, TokenTreeExt as _};
2use proc_macro::{Spacing::*, 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 { iter: self.into_iter(), buf: VecDeque::new() }
10    }
11}
12impl<I: IntoIterator<Item = TokenTree>> ParseIterExt for I { }
13
14/// Peek `n` iterator adapter
15#[derive(Debug, Clone)]
16pub struct ParseIter<I: Iterator<Item = TokenTree>> {
17    iter: I,
18    buf: VecDeque<TokenTree>,
19}
20
21impl<I: Iterator<Item = TokenTree>> ParseIter<I> {
22    pub fn peek(&mut self) -> Option<&TokenTree> {
23        self.peek_i(0)
24    }
25
26    pub fn next_if<F>(&mut self, f: F) -> Option<TokenTree>
27    where F: FnOnce(&TokenTree) -> bool,
28    {
29        let peek = self.peek()?;
30
31        if f(peek) {
32            self.next()
33        } else {
34            None
35        }
36    }
37
38    pub fn next_i_if<F>(&mut self, i: usize, f: F) -> Option<TokenTree>
39    where F: FnOnce(&TokenTree) -> bool,
40    {
41        let peek = self.peek_i(i)?;
42
43        if f(peek) {
44            self.nth(i)
45        } else {
46            None
47        }
48    }
49
50    /// # Panics
51    /// `self.count() < N`
52    #[track_caller]
53    pub fn next_tts<const N: usize>(&mut self) -> [TokenTree; N] {
54        array::from_fn(|_| self.next()
55            .expect("unexpected end of input"))
56    }
57
58    pub fn peek_i(&mut self, i: usize) -> Option<&TokenTree> {
59        for _ in self.buf.len()..=i {
60            self.buf.push_back(self.iter.next()?);
61        }
62        Some(&self.buf[i])
63    }
64
65    pub fn peek_is<F>(&mut self, f: F) -> bool
66    where F: FnOnce(&TokenTree) -> bool,
67    {
68        self.peek().is_some_and(f)
69    }
70
71    pub fn peek_i_is<F>(&mut self, i: usize, f: F) -> bool
72    where F: FnOnce(&TokenTree) -> bool,
73    {
74        self.peek_i(i).is_some_and(f)
75    }
76
77    /// Peek jointed puncts
78    pub fn peek_puncts(
79        &mut self,
80        puncts: impl AsRef<[u8]>,
81    ) -> Option<impl Iterator<Item = &TokenTree>> {
82        self.peek_i_puncts(0, puncts)
83    }
84
85    /// Peek jointed puncts
86    pub fn peek_i_puncts(
87        &mut self,
88        i: usize,
89        puncts: impl AsRef<[u8]>,
90    ) -> Option<impl Iterator<Item = &TokenTree>> {
91        let mut prev = None;
92
93        let puncts = puncts.as_ref();
94
95        for (j, ch) in puncts.iter().copied().map(char::from).enumerate() {
96            if let Some(prev) = prev {
97                if prev == Alone { return None }
98            }
99
100            let tt = self.peek_i(i + j)?;
101            let TokenTree::Punct(p) = tt else { return None };
102            if p.as_char() != ch { return None }
103
104            prev = Some(p.spacing());
105        }
106        Some(self.buf.iter().skip(i).take(puncts.len()))
107    }
108
109    /// Next jointed puncts
110    pub fn next_puncts(
111        &mut self,
112        puncts: impl AsRef<[u8]>,
113    ) -> Option<impl Iterator<Item = TokenTree> + '_> {
114        let puncts = puncts.as_ref();
115        let _ = self.peek_puncts(puncts)?;
116
117        Some(self.buf.drain(..puncts.len()))
118    }
119
120    /// Split [`TokenStream`] with `puncts`
121    ///
122    /// Like `"+-,-+".split_puncts(",")` -> `"+-"`
123    #[allow(clippy::missing_panics_doc)]
124    pub fn split_puncts(&mut self, puncts: impl AsRef<[u8]>) -> Option<TokenStream> {
125        let puncts = puncts.as_ref();
126        let mut i = 0;
127
128        loop {
129            if self.peek_i_puncts(i, puncts).is_some() {
130                let left = self.take(i).collect();
131                let _ = self.next_puncts(puncts).unwrap();
132                break Some(left);
133            }
134            self.peek_i(i)?;
135            i += 1;
136        }
137    }
138
139    /// Split all [`TokenStream`] with `puncts`
140    ///
141    /// Like `"+-,-+".split_puncts(",")` -> `"+-", "-+"`
142    pub fn split_puncts_all<'a>(
143        &'a mut self,
144        puncts: impl AsRef<[u8]> + 'a,
145    ) -> impl Iterator<Item = TokenStream> + 'a {
146        iter::from_fn(move || {
147            self.split_puncts(puncts.as_ref())
148                .or_else(|| self.peek().is_some().then(|| self.collect()))
149        })
150    }
151
152    pub fn next_attributes(&mut self) -> Vec<TokenStream> {
153        let mut attributes = vec![];
154
155        while self.peek_puncts("#").is_some()
156            &&self.peek_i_is(1, |tt| tt.is_delimiter_bracket())
157        {
158            attributes.push(self.next_tts::<2>().into_iter().collect());
159        }
160
161        attributes
162    }
163
164    pub fn next_outer_attributes(&mut self) -> Vec<TokenStream> {
165        let mut attributes = vec![];
166
167        while self.peek_puncts("#").is_some()
168            &&self.peek_i_is(1, |tt| tt.is_punch('!'))
169            &&self.peek_i_is(2, |tt| tt.is_delimiter_bracket())
170        {
171            attributes.push(stream(self.next_tts::<3>()));
172        }
173
174        attributes
175    }
176
177    #[allow(clippy::missing_panics_doc)]
178    pub fn next_vis(&mut self) -> Option<TokenStream> {
179        if self.peek_is(|tt| tt.is_keyword("pub")) {
180            if self.peek_i_is(1, |tt| tt.is_delimiter_paren()) {
181                return Some(stream(self.next_tts::<2>()));
182            }
183            return Some(self.next().unwrap().into());
184        }
185        None
186    }
187}
188impl<I: Iterator<Item = TokenTree>> Iterator for ParseIter<I> {
189    type Item = TokenTree;
190
191    fn next(&mut self) -> Option<Self::Item> {
192        self.buf.pop_front()
193            .or_else(|| self.iter.next())
194    }
195
196    fn count(self) -> usize
197    where Self: Sized,
198    {
199        self.buf.len() + self.iter.count()
200    }
201
202    fn fold<B, F>(self, init: B, f: F) -> B
203    where Self: Sized,
204          F: FnMut(B, Self::Item) -> B,
205    {
206        self.buf.into_iter().chain(self.iter).fold(init, f)
207    }
208
209    fn size_hint(&self) -> (usize, Option<usize>) {
210        let (lo, hi) = self.iter.size_hint();
211        let lo = lo.saturating_add(self.buf.len());
212        let hi = hi.and_then(|hi| hi.checked_add(self.buf.len()));
213        (lo, hi)
214    }
215}
216impl<I: DoubleEndedIterator<Item = TokenTree>> DoubleEndedIterator for ParseIter<I> {
217    fn next_back(&mut self) -> Option<Self::Item> {
218        self.iter.next_back()
219            .or_else(|| self.buf.pop_back())
220    }
221
222    fn rfold<B, F>(self, init: B, f: F) -> B
223    where Self: Sized,
224          F: FnMut(B, Self::Item) -> B,
225    {
226        self.buf.into_iter().chain(self.iter).rfold(init, f)
227    }
228}
229impl<I: ExactSizeIterator<Item = TokenTree>> ExactSizeIterator for ParseIter<I> { }
230impl<I: FusedIterator<Item = TokenTree>> FusedIterator for ParseIter<I> { }
proc_macro_tool/parse_iter.rs

proc_macro_tool/
parse_iter.rs
