easy_ext/
iter.rs

1// SPDX-License-Identifier: Apache-2.0 OR MIT
2
3use proc_macro::{
4    token_stream, Delimiter, Group, Ident, Literal, Punct, Spacing, TokenStream, TokenTree,
5};
6
7use crate::error::Result;
8
9#[derive(Clone)]
10pub(crate) struct TokenIter {
11    stack: Vec<token_stream::IntoIter>,
12    peeked: Option<TokenTree>,
13    peeked2: Option<TokenTree>,
14    peeked3: Option<TokenTree>,
15    peeked4: Option<TokenTree>,
16}
17
18impl TokenIter {
19    pub(crate) fn new(tokens: TokenStream) -> Self {
20        Self {
21            stack: vec![tokens.into_iter()],
22            peeked: None,
23            peeked2: None,
24            peeked3: None,
25            peeked4: None,
26        }
27    }
28
29    #[allow(clippy::wrong_self_convention)]
30    pub(crate) fn is_empty(&mut self) -> bool {
31        self.peek().is_none()
32    }
33
34    pub(crate) fn peek(&mut self) -> Option<&TokenTree> {
35        self.peeked = self.next();
36        self.peeked.as_ref()
37    }
38
39    pub(crate) fn peek_t(&mut self, token: &dyn Token) -> bool {
40        match self.peek() {
41            Some(tt) => token.match_token(tt),
42            None => false,
43        }
44    }
45
46    pub(crate) fn peek2(&mut self) -> Option<&TokenTree> {
47        let peeked = self.next();
48        let peeked2 = self.next();
49        self.peeked = peeked;
50        self.peeked2 = peeked2;
51        self.peeked2.as_ref()
52    }
53
54    pub(crate) fn peek2_t(&mut self, token: &dyn Token) -> bool {
55        match self.peek2() {
56            Some(tt) => token.match_token(tt),
57            None => false,
58        }
59    }
60
61    pub(crate) fn peek3(&mut self) -> Option<&TokenTree> {
62        let peeked = self.next();
63        let peeked2 = self.next();
64        let peeked3 = self.next();
65        self.peeked = peeked;
66        self.peeked2 = peeked2;
67        self.peeked3 = peeked3;
68        self.peeked3.as_ref()
69    }
70
71    pub(crate) fn peek3_t(&mut self, token: &dyn Token) -> bool {
72        match self.peek3() {
73            Some(tt) => token.match_token(tt),
74            None => false,
75        }
76    }
77
78    pub(crate) fn peek4(&mut self) -> Option<&TokenTree> {
79        let peeked = self.next();
80        let peeked2 = self.next();
81        let peeked3 = self.next();
82        let peeked4 = self.next();
83        self.peeked = peeked;
84        self.peeked2 = peeked2;
85        self.peeked3 = peeked3;
86        self.peeked4 = peeked4;
87        self.peeked4.as_ref()
88    }
89
90    pub(crate) fn peek4_t(&mut self, token: &dyn Token) -> bool {
91        match self.peek4() {
92            Some(tt) => token.match_token(tt),
93            None => false,
94        }
95    }
96
97    pub(crate) fn peek_ident(&mut self) -> Option<&Ident> {
98        match self.peek() {
99            Some(TokenTree::Ident(i)) => Some(i),
100            _ => None,
101        }
102    }
103
104    pub(crate) fn peek2_ident(&mut self) -> Option<&Ident> {
105        match self.peek2() {
106            Some(TokenTree::Ident(i)) => Some(i),
107            _ => None,
108        }
109    }
110
111    pub(crate) fn peek3_ident(&mut self) -> Option<&Ident> {
112        match self.peek3() {
113            Some(TokenTree::Ident(i)) => Some(i),
114            _ => None,
115        }
116    }
117
118    pub(crate) fn parse_ident(&mut self) -> Result<Ident> {
119        match self.next() {
120            Some(TokenTree::Ident(i)) => Ok(i),
121            // TODO: pass scope span if tt is None
122            tt => bail!(tt, "expected identifier"),
123        }
124    }
125
126    pub(crate) fn parse_ident_opt(&mut self) -> Option<Ident> {
127        self.peek_ident()?;
128        Some(self.parse_ident().unwrap())
129    }
130
131    pub(crate) fn parse_kw(&mut self, kw: &str) -> Result<Ident> {
132        let tt = self.next();
133        match &tt {
134            Some(TokenTree::Ident(i)) if i.to_string() == kw => {
135                if let Some(TokenTree::Ident(i)) = tt {
136                    Ok(i)
137                } else {
138                    unreachable!()
139                }
140            }
141            // TODO: pass scope span if tt is None
142            tt => bail!(tt, "expected `{}`", kw),
143        }
144    }
145
146    pub(crate) fn parse_kw_opt(&mut self, kw: &str) -> Option<Ident> {
147        if self.peek_t(&kw) {
148            Some(self.parse_ident().unwrap())
149        } else {
150            None
151        }
152    }
153
154    pub(crate) fn peek_punct(&mut self, ch: char) -> Option<&Punct> {
155        match self.peek() {
156            Some(TokenTree::Punct(p)) if p.as_char() == ch => Some(p),
157            _ => None,
158        }
159    }
160
161    pub(crate) fn peek2_punct(&mut self, ch: char) -> Option<&Punct> {
162        match self.peek2() {
163            Some(TokenTree::Punct(p)) if p.as_char() == ch => Some(p),
164            _ => None,
165        }
166    }
167
168    pub(crate) fn parse_punct(&mut self, ch: char) -> Result<Punct> {
169        let tt = self.next();
170        match &tt {
171            Some(TokenTree::Punct(p)) if p.as_char() == ch => {
172                if let Some(TokenTree::Punct(p)) = tt {
173                    Ok(p)
174                } else {
175                    unreachable!()
176                }
177            }
178            // TODO: pass scope span if tt is None
179            tt => bail!(tt, "expected `{}`", ch),
180        }
181    }
182
183    pub(crate) fn parse_punct_opt(&mut self, ch: char) -> Option<Punct> {
184        self.peek_punct(ch)?;
185        Some(self.parse_punct(ch).unwrap())
186    }
187
188    pub(crate) fn peek_lifetime(&mut self) -> bool {
189        self.peek_punct('\'').map_or(false, |p| p.spacing() == Spacing::Joint)
190            && self.peek2_ident().is_some()
191    }
192
193    pub(crate) fn peek2_lifetime(&mut self) -> bool {
194        self.peek2_punct('\'').map_or(false, |p| p.spacing() == Spacing::Joint)
195            && self.peek3_ident().is_some()
196    }
197
198    pub(crate) fn parse_group(&mut self, delimiter: Delimiter) -> Result<Group> {
199        let tt = self.next();
200        match &tt {
201            Some(TokenTree::Group(g)) if g.delimiter() == delimiter => {
202                if let Some(TokenTree::Group(g)) = tt {
203                    Ok(g)
204                } else {
205                    unreachable!()
206                }
207            }
208            tt => {
209                let d = match delimiter {
210                    Delimiter::Brace => "`{`",
211                    Delimiter::Bracket => "`[`",
212                    Delimiter::Parenthesis => "`(`",
213                    Delimiter::None => "none-delimited group",
214                };
215                // TODO: pass scope span if tt is None
216                bail!(tt, "expected {}", d)
217            }
218        }
219    }
220
221    pub(crate) fn peek_literal(&mut self) -> Option<&Literal> {
222        match self.peek() {
223            Some(TokenTree::Literal(l)) => Some(l),
224            _ => None,
225        }
226    }
227
228    pub(crate) fn parse_literal(&mut self) -> Result<Literal> {
229        match self.next() {
230            Some(TokenTree::Literal(l)) => Ok(l),
231            // TODO: pass scope span if tt is None
232            tt => bail!(tt, "expected literal"),
233        }
234    }
235
236    pub(crate) fn parse_literal_opt(&mut self) -> Option<Literal> {
237        self.peek_literal()?;
238        Some(self.parse_literal().unwrap())
239    }
240
241    pub(crate) fn tt(&mut self) -> Result<TokenTree> {
242        self.next().ok_or_else(|| {
243            // TODO: pass scope span
244            format_err!(TokenStream::new(), "unexpected end of input")
245        })
246    }
247}
248
249// Based on https://github.com/dtolnay/proc-macro-hack/blob/0.5.19/src/iter.rs
250impl Iterator for TokenIter {
251    type Item = TokenTree;
252
253    fn next(&mut self) -> Option<Self::Item> {
254        if let Some(tt) = self.peeked.take() {
255            return Some(tt);
256        }
257        if let Some(tt) = self.peeked2.take() {
258            return Some(tt);
259        }
260        if let Some(tt) = self.peeked3.take() {
261            return Some(tt);
262        }
263        if let Some(tt) = self.peeked4.take() {
264            return Some(tt);
265        }
266        loop {
267            let top = self.stack.last_mut()?;
268            match top.next() {
269                None => drop(self.stack.pop()),
270                Some(TokenTree::Group(ref group)) if group.delimiter() == Delimiter::None => {
271                    self.stack.push(group.stream().into_iter());
272                }
273                Some(tt) => return Some(tt),
274            }
275        }
276    }
277}
278
279pub(crate) trait Token {
280    fn match_token(&self, tt: &TokenTree) -> bool;
281}
282
283impl Token for char {
284    fn match_token(&self, tt: &TokenTree) -> bool {
285        match tt {
286            TokenTree::Punct(p) => p.as_char() == *self,
287            _ => false,
288        }
289    }
290}
291
292impl Token for &str {
293    fn match_token(&self, tt: &TokenTree) -> bool {
294        match tt {
295            TokenTree::Ident(i) => i.to_string() == *self,
296            _ => false,
297        }
298    }
299}
300
301impl Token for Delimiter {
302    fn match_token(&self, tt: &TokenTree) -> bool {
303        match tt {
304            TokenTree::Group(g) => g.delimiter() == *self,
305            _ => false,
306        }
307    }
308}
easy_ext/iter.rs

easy_ext/
iter.rs
