proc_macro_tool/
exts.rs

1use core::{iter::once, mem::take};
2
3use crate::{
4    puncts, puncts_spanned, ParseIter, ParseIterExt as _, SetSpan as _,
5    TokenKind,
6};
7use proc_macro::{
8    Delimiter, Group, Ident, Literal, Punct,
9    Spacing::{self, *},
10    Span, TokenStream, TokenTree,
11};
12
13pub trait TokenStreamExt
14    : Default
15    + Extend<TokenTree>
16    + Extend<TokenStream>
17    + IntoIterator<Item = TokenTree>
18    + Sized
19{
20    /// Extend a [`TokenTree`]
21    fn push(&mut self, tt: TokenTree) -> &mut Self {
22        self.extend(once(tt));
23        self
24    }
25
26    /// Extend a [`TokenStream`]
27    fn add(&mut self, stream: TokenStream) -> &mut Self {
28        self.extend(once(stream));
29        self
30    }
31
32    /// Call [`mem::take`](std::mem::take)
33    #[must_use]
34    fn take(&mut self) -> Self {
35        take(self)
36    }
37
38    /// Call [`Group::new`]
39    fn grouped(self, delimiter: Delimiter) -> Group;
40
41    fn grouped_paren(self) -> Group {
42        self.grouped(Delimiter::Parenthesis)
43    }
44
45    fn grouped_brace(self) -> Group {
46        self.grouped(Delimiter::Brace)
47    }
48
49    fn grouped_bracket(self) -> Group {
50        self.grouped(Delimiter::Bracket)
51    }
52
53    fn grouped_none(self) -> Group {
54        self.grouped(Delimiter::None)
55    }
56
57    /// Split [`TokenStream`] to `predicate` false and true
58    ///
59    /// Like `"+-,-+".split_puncts(",")` -> `("+-", "-+")`
60    fn split_puncts(self, puncts: impl AsRef<[u8]>) -> Option<(
61        Self,
62        ParseIter<Self::IntoIter>,
63    )>;
64}
65impl TokenStreamExt for TokenStream {
66    fn grouped(self, delimiter: Delimiter) -> Group {
67        Group::new(delimiter, self)
68    }
69
70    fn split_puncts(self, puncts: impl AsRef<[u8]>) -> Option<(
71        Self,
72        ParseIter<Self::IntoIter>,
73    )>
74    {
75        let mut iter = self.parse_iter();
76        Some((iter.split_puncts(puncts)?, iter))
77    }
78
79}
80
81/// Remake each subtree methods
82pub trait WalkExt
83    : IntoIterator<Item = TokenTree>
84    + FromIterator<TokenTree>
85{
86    /// Remake each subtree
87    ///
88    /// `"(1+2)*3"` -> call `f` on `1`, `+`, `2`, `(f(1) f(+) f(2))`, `*`, `3`
89    #[must_use]
90    fn walk<F>(self, mut f: F) -> Self
91    where F: FnMut(TokenTree) -> TokenTree
92    {
93        fn walk_impl<I, F>(this: I, f: &mut F) -> I
94        where I: IntoIterator<Item = TokenTree> + FromIterator<TokenTree>,
95              F: FnMut(TokenTree) -> TokenTree
96        {
97            this.into_iter()
98                .map(|tt| {
99                    let tt = match tt {
100                        TokenTree::Group(g) => {
101                            walk_impl(g.stream(), &mut *f)
102                                .grouped(g.delimiter())
103                                .set_spaned(g.span())
104                                .into()
105                        },
106                        _ => tt,
107                    };
108                    f(tt)
109                })
110                .collect()
111        }
112        walk_impl(self, &mut f)
113    }
114}
115impl<I: IntoIterator<Item = TokenTree> + FromIterator<TokenTree>> WalkExt for I { }
116
117pub trait TokenTreeExt: Into<TokenTree> + Sized {
118    fn as_ident(&self) -> Option<&Ident>     { None }
119    fn as_punct(&self) -> Option<&Punct>     { None }
120    fn as_group(&self) -> Option<&Group>     { None }
121    fn as_literal(&self) -> Option<&Literal> { None }
122    fn into_ident(self) -> Result<Ident, Self>     { Err(self) }
123    fn into_punct(self) -> Result<Punct, Self>     { Err(self) }
124    fn into_group(self) -> Result<Group, Self>     { Err(self) }
125    fn into_literal(self) -> Result<Literal, Self> { Err(self) }
126
127    fn to_ident(&self) -> Result<&Ident, &Self> {
128        self.as_ident().ok_or(self)
129    }
130
131    fn to_punct(&self) -> Result<&Punct, &Self> {
132        self.as_punct().ok_or(self)
133    }
134
135    fn to_group(&self) -> Result<&Group, &Self> {
136        self.as_group().ok_or(self)
137    }
138
139    fn to_literal(&self) -> Result<&Literal, &Self> {
140        self.as_literal().ok_or(self)
141    }
142
143    fn is_ident(&self) -> bool {
144        self.as_ident().is_some()
145    }
146
147    fn is_punct(&self) -> bool {
148        self.as_punct().is_some()
149    }
150
151    fn is_group(&self) -> bool {
152        self.as_group().is_some()
153    }
154
155    fn is_literal(&self) -> bool {
156        self.as_literal().is_some()
157    }
158
159    fn kind(&self) -> TokenKind {
160        if self.is_literal() {
161            TokenKind::Literal
162        } else if self.is_punct() {
163            TokenKind::Punct
164        } else if self.is_group() {
165            TokenKind::Group
166        } else if self.is_ident() {
167            TokenKind::Ident
168        } else {
169            unimplemented!()
170        }
171    }
172
173    /// Ident content equal to `keyword` str
174    ///
175    /// Other return `false` when `self` is not [`Ident`]
176    fn is_keyword(&self, keyword: &str) -> bool {
177        self.as_ident().is_some_and(|i| i.to_string() == keyword)
178    }
179
180    /// Punct char equal to `ch`
181    ///
182    /// Other return `false` when `self` is not [`Punct`]
183    fn is_punch(&self, ch: char) -> bool {
184        self.as_punct().is_some_and(|p| p.as_char() == ch)
185    }
186
187    /// Group delimiter is not [`Delimiter::None`]
188    ///
189    /// Other return `false` when `self` is not [`Group`]
190    fn is_solid_group(&self) -> bool {
191        self.as_group().is_some_and(|g| g.is_solid_group())
192    }
193
194    /// Group delimiter equal to `delimiter`
195    ///
196    /// Other return `false` when `self` is not [`Group`]
197    fn is_delimiter(&self, delimiter: Delimiter) -> bool {
198        self.as_group().is_some_and(|g| g.is_delimiter(delimiter))
199    }
200
201    /// Punct spacing is [`Joint`]
202    ///
203    /// Other return `false` when `self` is not [`Punct`]
204    fn is_joint(&self) -> bool {
205        self.as_punct().is_some_and(|p| p.spacing() == Joint)
206    }
207
208    fn as_punct_char(&self) -> Option<char> {
209        self.as_punct().map(|p| p.as_char())
210    }
211
212    /// [`Into`] [`TokenTree`], like [`TokenTree::from(self)`]
213    ///
214    /// [`TokenTree::from(self)`]: TokenTree::from
215    fn tt(self) -> TokenTree {
216        self.into()
217    }
218
219    /// [`TokenStream::from_iter(self.tt())`](TokenStream::from_iter)
220    fn unit_stream(self) -> TokenStream {
221        self.tt().into()
222    }
223}
224impl TokenTreeExt for TokenTree {
225    fn as_ident(&self) -> Option<&Ident> {
226        match self {
227            TokenTree::Ident(i) => Some(i),
228            _ => None,
229        }
230    }
231
232    fn as_punct(&self) -> Option<&Punct> {
233        match self {
234            TokenTree::Punct(i) => Some(i),
235            _ => None,
236        }
237    }
238
239    fn as_group(&self) -> Option<&Group> {
240        match self {
241            TokenTree::Group(i) => Some(i),
242            _ => None,
243        }
244    }
245
246    fn as_literal(&self) -> Option<&Literal> {
247        match self {
248            TokenTree::Literal(i) => Some(i),
249            _ => None,
250        }
251    }
252
253    fn into_ident(self) -> Result<Ident, Self> {
254        match self {
255            TokenTree::Ident(i) => Ok(i),
256            _ => Err(self),
257        }
258    }
259
260    fn into_punct(self) -> Result<Punct, Self> {
261        match self {
262            TokenTree::Punct(i) => Ok(i),
263            _ => Err(self),
264        }
265    }
266
267    fn into_group(self) -> Result<Group, Self> {
268        match self {
269            TokenTree::Group(i) => Ok(i),
270            _ => Err(self),
271        }
272    }
273
274    fn into_literal(self) -> Result<Literal, Self> {
275        match self {
276            TokenTree::Literal(i) => Ok(i),
277            _ => Err(self),
278        }
279    }
280
281    fn kind(&self) -> TokenKind {
282        self.into()
283    }
284}
285macro_rules! impl_token_tree_ext {
286    ($as:ident, $into:ident, $ty:ident) => {
287        impl TokenTreeExt for $ty {
288            fn $as(&self) -> Option<&$ty> {
289                Some(self)
290            }
291            fn $into(self) -> Result<$ty, Self> {
292                Ok(self)
293            }
294            fn kind(&self) -> TokenKind {
295                TokenKind::$ty
296            }
297        }
298    };
299}
300impl_token_tree_ext!(as_ident,   into_ident,   Ident);
301impl_token_tree_ext!(as_punct,   into_punct,   Punct);
302impl_token_tree_ext!(as_literal, into_literal, Literal);
303impl TokenTreeExt for Group {
304    fn as_group(&self) -> Option<&Group> {
305        Some(self)
306    }
307
308    fn into_group(self) -> Result<Group, Self> {
309        Ok(self)
310    }
311
312    fn kind(&self) -> TokenKind {
313        TokenKind::Group
314    }
315
316    fn is_solid_group(&self) -> bool {
317        self.delimiter() != Delimiter::None
318    }
319
320    fn is_delimiter(&self, delimiter: Delimiter) -> bool {
321        self.delimiter() == delimiter
322    }
323}
324
325/// Create unsuffixed [`Literal`]
326pub trait Unsuffixed {
327    fn unsuffixed(self) -> Literal;
328}
329/// Create suffixed [`Literal`]
330pub trait Suffixed {
331    fn suffixed(self) -> Literal;
332}
333macro_rules! impl_unsuffixes {
334    ( $($ty:ty: $unsuffixed:ident $($suffixed:ident)?);+ $(;)? ) => {
335        $(
336            #[doc = concat!(
337                "Call [`Literal::",
338                stringify!($unsuffixed),
339                "`]",
340            )]
341            impl Unsuffixed for $ty {
342                fn unsuffixed(self) -> Literal {
343                    Literal::$unsuffixed(self)
344                }
345            }
346
347            $(
348                #[doc = concat!(
349                    "Call [`Literal::",
350                    stringify!($unsuffixed),
351                    "`]",
352                )]
353                impl Suffixed for $ty {
354                    fn suffixed(self) -> Literal {
355                        Literal::$suffixed(self)
356                    }
357                }
358            )?
359        )*
360    };
361}
362impl_unsuffixes! {
363    i8:     i8_unsuffixed       i8_suffixed;
364    i16:    i16_unsuffixed      i16_suffixed;
365    i32:    i32_unsuffixed      i32_suffixed;
366    i64:    i64_unsuffixed      i64_suffixed;
367    i128:   i128_unsuffixed     i128_suffixed;
368    u8:     u8_unsuffixed       u8_suffixed;
369    u16:    u16_unsuffixed      u16_suffixed;
370    u32:    u32_unsuffixed      u32_suffixed;
371    u64:    u64_unsuffixed      u64_suffixed;
372    u128:   u128_unsuffixed     u128_suffixed;
373    f32:    f32_unsuffixed      f32_suffixed;
374    f64:    f64_unsuffixed      f64_suffixed;
375    usize:  usize_unsuffixed    usize_suffixed;
376    isize:  isize_unsuffixed    isize_suffixed;
377    char:   character;
378    &str:   string;
379    &[u8]:  byte_string;
380}
381
382pub trait PunctsExt: AsRef<[u8]> {
383    /// Call [`puncts`]
384    fn puncts(&self) -> TokenStream {
385        puncts(self)
386    }
387
388    /// Call [`puncts_spanned`]
389    fn puncts_spanned(&self, span: Span) -> TokenStream {
390        puncts_spanned(self, span)
391    }
392}
393impl<T: AsRef<[u8]> + ?Sized> PunctsExt for T { }
394
395pub trait PunctExt {
396    fn punct(self, spacing: Spacing) -> Punct;
397}
398impl PunctExt for char {
399    /// Call [`Punct::new`]
400    fn punct(self, spacing: Spacing) -> Punct {
401        Punct::new(self, spacing)
402    }
403}
404
405pub trait StrExt {
406    fn ident(&self, span: Span) -> Ident;
407}
408impl StrExt for str {
409    /// Call [`Ident::new`]
410    fn ident(&self, span: Span) -> Ident {
411        Ident::new(self, span)
412    }
413}
proc_macro_tool/exts.rs

proc_macro_tool/
exts.rs
