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                            g.map(|this| walk_impl(this, f)).tt()
102                        },
103                        _ => tt,
104                    };
105                    f(tt)
106                })
107                .collect()
108        }
109        walk_impl(self, &mut f)
110    }
111}
112impl<I: IntoIterator<Item = TokenTree> + FromIterator<TokenTree>> WalkExt for I { }
113
114pub trait TokenTreeExt: Into<TokenTree> + Sized {
115    fn as_ident(&self) -> Option<&Ident>     { None }
116    fn as_punct(&self) -> Option<&Punct>     { None }
117    fn as_group(&self) -> Option<&Group>     { None }
118    fn as_literal(&self) -> Option<&Literal> { None }
119    fn into_ident(self) -> Result<Ident, Self>     { Err(self) }
120    fn into_punct(self) -> Result<Punct, Self>     { Err(self) }
121    fn into_group(self) -> Result<Group, Self>     { Err(self) }
122    fn into_literal(self) -> Result<Literal, Self> { Err(self) }
123
124    fn to_ident(&self) -> Result<&Ident, &Self> {
125        self.as_ident().ok_or(self)
126    }
127
128    fn to_punct(&self) -> Result<&Punct, &Self> {
129        self.as_punct().ok_or(self)
130    }
131
132    fn to_group(&self) -> Result<&Group, &Self> {
133        self.as_group().ok_or(self)
134    }
135
136    fn to_literal(&self) -> Result<&Literal, &Self> {
137        self.as_literal().ok_or(self)
138    }
139
140    fn is_ident(&self) -> bool {
141        self.as_ident().is_some()
142    }
143
144    fn is_punct(&self) -> bool {
145        self.as_punct().is_some()
146    }
147
148    fn is_group(&self) -> bool {
149        self.as_group().is_some()
150    }
151
152    fn is_literal(&self) -> bool {
153        self.as_literal().is_some()
154    }
155
156    fn kind(&self) -> TokenKind {
157        if self.is_literal() {
158            TokenKind::Literal
159        } else if self.is_punct() {
160            TokenKind::Punct
161        } else if self.is_group() {
162            TokenKind::Group
163        } else if self.is_ident() {
164            TokenKind::Ident
165        } else {
166            unimplemented!()
167        }
168    }
169
170    /// Ident content equal to `keyword` str
171    ///
172    /// Other return `false` when `self` is not [`Ident`]
173    fn is_keyword(&self, keyword: &str) -> bool {
174        self.as_ident().is_some_and(|i| i.to_string() == keyword)
175    }
176
177    /// Punct char equal to `ch`
178    ///
179    /// Other return `false` when `self` is not [`Punct`]
180    fn is_punch(&self, ch: char) -> bool {
181        self.as_punct().is_some_and(|p| p.as_char() == ch)
182    }
183
184    /// Group delimiter is not [`Delimiter::None`]
185    ///
186    /// Other return `false` when `self` is not [`Group`]
187    fn is_solid_group(&self) -> bool {
188        self.as_group().is_some_and(|g| g.is_solid_group())
189    }
190
191    /// Group delimiter equal to `delimiter`
192    ///
193    /// Other return `false` when `self` is not [`Group`]
194    fn is_delimiter(&self, delimiter: Delimiter) -> bool {
195        self.as_group().is_some_and(|g| g.is_delimiter(delimiter))
196    }
197
198    /// Like [`self.is_delimiter(Delimiter::Parenthesis)`](#method.is_delimiter)
199    fn is_delimiter_paren(&self) -> bool {
200        self.is_delimiter(Delimiter::Parenthesis)
201    }
202
203    /// Like [`self.is_delimiter(Delimiter::Brace)`](#method.is_delimiter)
204    fn is_delimiter_brace(&self) -> bool {
205        self.is_delimiter(Delimiter::Brace)
206    }
207
208    /// Like [`self.is_delimiter(Delimiter::Bracket)`](#method.is_delimiter)
209    fn is_delimiter_bracket(&self) -> bool {
210        self.is_delimiter(Delimiter::Bracket)
211    }
212
213    /// Like [`self.is_delimiter(Delimiter::None)`](#method.is_delimiter)
214    fn is_delimiter_none(&self) -> bool {
215        self.is_delimiter(Delimiter::None)
216    }
217
218    /// Call [`Group::stream`] when [`delimiter`] is [`Delimiter::Parenthesis`]
219    ///
220    /// Other return `false` when `self` is not [`Group`]
221    ///
222    /// [`delimiter`]: Group::delimiter
223    fn to_paren_stream(&self) -> Result<TokenStream, &Self> {
224        self.as_group()
225            .and_then(|g| g.is_delimiter_paren().then(|| g.stream()))
226            .ok_or(self)
227    }
228
229    /// Call [`Group::stream`] when [`delimiter`] is [`Delimiter::Brace`]
230    ///
231    /// Other return `false` when `self` is not [`Group`]
232    ///
233    /// [`delimiter`]: Group::delimiter
234    fn to_brace_stream(&self) -> Result<TokenStream, &Self> {
235        self.as_group()
236            .and_then(|g| g.is_delimiter_brace().then(|| g.stream()))
237            .ok_or(self)
238    }
239
240    /// Call [`Group::stream`] when [`delimiter`] is [`Delimiter::Bracket`]
241    ///
242    /// Other return `false` when `self` is not [`Group`]
243    ///
244    /// [`delimiter`]: Group::delimiter
245    fn to_bracket_stream(&self) -> Result<TokenStream, &Self> {
246        self.as_group()
247            .and_then(|g| g.is_delimiter_bracket().then(|| g.stream()))
248            .ok_or(self)
249    }
250
251    /// Call [`Group::stream`] when [`delimiter`] is [`Delimiter::None`]
252    ///
253    /// Other return `false` when `self` is not [`Group`]
254    ///
255    /// [`delimiter`]: Group::delimiter
256    fn to_none_stream(&self) -> Result<TokenStream, &Self> {
257        self.as_group()
258            .and_then(|g| g.is_delimiter_none().then(|| g.stream()))
259            .ok_or(self)
260    }
261
262    /// Like [`to_paren_stream`](#method.to_paren_stream),
263    /// but using `Self` instead of `&Self`
264    fn into_paren_stream(self) -> Result<TokenStream, Self> {
265        self.to_paren_stream()
266            .ok().ok_or(self)
267    }
268
269    /// Like [`to_brace_stream`](#method.to_brace_stream),
270    /// but using `Self` instead of `&Self`
271    fn into_brace_stream(self) -> Result<TokenStream, Self> {
272        self.to_brace_stream()
273            .ok().ok_or(self)
274    }
275
276    /// Like [`to_bracket_stream`](#method.to_bracket_stream),
277    /// but using `Self` instead of `&Self`
278    fn into_bracket_stream(self) -> Result<TokenStream, Self> {
279        self.to_bracket_stream()
280            .ok().ok_or(self)
281    }
282
283    /// Like [`to_none_stream`](#method.to_none_stream),
284    /// but using `Self` instead of `&Self`
285    fn into_none_stream(self) -> Result<TokenStream, Self> {
286        self.to_none_stream()
287            .ok().ok_or(self)
288    }
289
290    /// Punct spacing is [`Joint`]
291    ///
292    /// Other return `false` when `self` is not [`Punct`]
293    fn is_joint(&self) -> bool {
294        self.as_punct().is_some_and(|p| p.spacing() == Joint)
295    }
296
297    fn as_punct_char(&self) -> Option<char> {
298        self.as_punct().map(|p| p.as_char())
299    }
300
301    /// [`Into`] [`TokenTree`], like [`TokenTree::from(self)`]
302    ///
303    /// [`TokenTree::from(self)`]: TokenTree::from
304    fn tt(self) -> TokenTree {
305        self.into()
306    }
307
308    /// [`TokenStream::from_iter(self.tt())`](TokenStream::from_iter)
309    fn unit_stream(self) -> TokenStream {
310        self.tt().into()
311    }
312}
313impl TokenTreeExt for TokenTree {
314    fn as_ident(&self) -> Option<&Ident> {
315        match self {
316            TokenTree::Ident(i) => Some(i),
317            _ => None,
318        }
319    }
320
321    fn as_punct(&self) -> Option<&Punct> {
322        match self {
323            TokenTree::Punct(i) => Some(i),
324            _ => None,
325        }
326    }
327
328    fn as_group(&self) -> Option<&Group> {
329        match self {
330            TokenTree::Group(i) => Some(i),
331            _ => None,
332        }
333    }
334
335    fn as_literal(&self) -> Option<&Literal> {
336        match self {
337            TokenTree::Literal(i) => Some(i),
338            _ => None,
339        }
340    }
341
342    fn into_ident(self) -> Result<Ident, Self> {
343        match self {
344            TokenTree::Ident(i) => Ok(i),
345            _ => Err(self),
346        }
347    }
348
349    fn into_punct(self) -> Result<Punct, Self> {
350        match self {
351            TokenTree::Punct(i) => Ok(i),
352            _ => Err(self),
353        }
354    }
355
356    fn into_group(self) -> Result<Group, Self> {
357        match self {
358            TokenTree::Group(i) => Ok(i),
359            _ => Err(self),
360        }
361    }
362
363    fn into_literal(self) -> Result<Literal, Self> {
364        match self {
365            TokenTree::Literal(i) => Ok(i),
366            _ => Err(self),
367        }
368    }
369
370    fn kind(&self) -> TokenKind {
371        self.into()
372    }
373}
374macro_rules! impl_token_tree_ext {
375    ($as:ident, $into:ident, $ty:ident) => {
376        impl TokenTreeExt for $ty {
377            fn $as(&self) -> Option<&$ty> {
378                Some(self)
379            }
380            fn $into(self) -> Result<$ty, Self> {
381                Ok(self)
382            }
383            fn kind(&self) -> TokenKind {
384                TokenKind::$ty
385            }
386        }
387    };
388}
389impl_token_tree_ext!(as_ident,   into_ident,   Ident);
390impl_token_tree_ext!(as_punct,   into_punct,   Punct);
391impl_token_tree_ext!(as_literal, into_literal, Literal);
392impl TokenTreeExt for Group {
393    fn as_group(&self) -> Option<&Group> {
394        Some(self)
395    }
396
397    fn into_group(self) -> Result<Group, Self> {
398        Ok(self)
399    }
400
401    fn kind(&self) -> TokenKind {
402        TokenKind::Group
403    }
404
405    fn is_solid_group(&self) -> bool {
406        self.delimiter() != Delimiter::None
407    }
408
409    fn is_delimiter(&self, delimiter: Delimiter) -> bool {
410        self.delimiter() == delimiter
411    }
412}
413
414/// Create unsuffixed [`Literal`]
415pub trait Unsuffixed {
416    fn unsuffixed(self) -> Literal;
417}
418/// Create suffixed [`Literal`]
419pub trait Suffixed {
420    fn suffixed(self) -> Literal;
421}
422macro_rules! impl_unsuffixes {
423    ( $($ty:ty: $unsuffixed:ident $($suffixed:ident)?);+ $(;)? ) => {
424        $(
425            #[doc = concat!(
426                "Call [`Literal::",
427                stringify!($unsuffixed),
428                "`]",
429            )]
430            impl Unsuffixed for $ty {
431                fn unsuffixed(self) -> Literal {
432                    Literal::$unsuffixed(self)
433                }
434            }
435
436            $(
437                #[doc = concat!(
438                    "Call [`Literal::",
439                    stringify!($suffixed),
440                    "`]",
441                )]
442                impl Suffixed for $ty {
443                    fn suffixed(self) -> Literal {
444                        Literal::$suffixed(self)
445                    }
446                }
447            )?
448        )*
449    };
450}
451impl_unsuffixes! {
452    i8:     i8_unsuffixed       i8_suffixed;
453    i16:    i16_unsuffixed      i16_suffixed;
454    i32:    i32_unsuffixed      i32_suffixed;
455    i64:    i64_unsuffixed      i64_suffixed;
456    i128:   i128_unsuffixed     i128_suffixed;
457    u8:     u8_unsuffixed       u8_suffixed;
458    u16:    u16_unsuffixed      u16_suffixed;
459    u32:    u32_unsuffixed      u32_suffixed;
460    u64:    u64_unsuffixed      u64_suffixed;
461    u128:   u128_unsuffixed     u128_suffixed;
462    f32:    f32_unsuffixed      f32_suffixed;
463    f64:    f64_unsuffixed      f64_suffixed;
464    usize:  usize_unsuffixed    usize_suffixed;
465    isize:  isize_unsuffixed    isize_suffixed;
466    char:   character;
467    &str:   string;
468    &[u8]:  byte_string;
469}
470
471pub trait PunctsExt: AsRef<[u8]> {
472    /// Call [`puncts`]
473    fn puncts(&self) -> TokenStream {
474        puncts(self)
475    }
476
477    /// Call [`puncts_spanned`]
478    fn puncts_spanned(&self, span: Span) -> TokenStream {
479        puncts_spanned(self, span)
480    }
481}
482impl<T: AsRef<[u8]> + ?Sized> PunctsExt for T { }
483
484pub trait PunctExt: Sized {
485    fn punct(self, spacing: Spacing) -> Punct;
486
487    /// Like [`.punct(Joint)`](#method.punct)
488    fn joint(self) -> Punct {
489        self.punct(Joint)
490    }
491
492    /// Like [`.punct(Alone)`](#method.punct)
493    fn alone(self) -> Punct {
494        self.punct(Alone)
495    }
496}
497impl PunctExt for char {
498    /// Call [`Punct::new`]
499    fn punct(self, spacing: Spacing) -> Punct {
500        Punct::new(self, spacing)
501    }
502}
503
504pub trait StrExt {
505    fn ident(&self, span: Span) -> Ident;
506}
507impl StrExt for str {
508    /// Call [`Ident::new`]
509    fn ident(&self, span: Span) -> Ident {
510        Ident::new(self, span)
511    }
512}
513
514pub trait GroupExt {
515    #[must_use]
516    fn map<F>(&self, f: F) -> Self
517    where F: FnOnce(TokenStream) -> TokenStream;
518}
519impl GroupExt for Group {
520    fn map<F>(&self, f: F) -> Self
521    where F: FnOnce(TokenStream) -> TokenStream,
522    {
523        f(self.stream())
524            .grouped(self.delimiter())
525            .set_spaned(self.span())
526    }
527}
proc_macro_tool/exts.rs

proc_macro_tool/
exts.rs
