proc_macro_tool/
exts.rs

1use core::{iter::once, mem::take};
2use std::convert::identity;
3
4use crate::{
5    puncts, puncts_spanned, ParseIter, ParseIterExt as _, SetSpan as _,
6    TokenKind,
7};
8use proc_macro::{
9    Delimiter, Group, Ident, Literal, Punct,
10    Spacing::{self, *},
11    Span, TokenStream, TokenTree,
12};
13
14pub trait StreamIterExt: Iterator + Sized {
15    fn join<I>(self, sep: I) -> TokenStream
16    where I: Clone,
17          TokenStream: Extend<I> + Extend<Self::Item>,
18    {
19        self.map(Ok::<_, ()>)
20            .try_join(sep)
21            .unwrap()
22    }
23
24    fn try_join<I, T, E>(mut self, sep: I) -> Result<TokenStream, E>
25    where Self: Iterator<Item = Result<T, E>>,
26          I: Clone,
27          TokenStream: Extend<I> + Extend<T>,
28    {
29        let mut result = TokenStream::new();
30
31        if let Some(first) = self.next() {
32            result.extend(once(first?));
33
34            for rest in self {
35                result.extend(once(sep.clone()));
36                result.extend(once(rest?));
37            }
38        }
39
40        Ok(result)
41    }
42}
43impl<T: Iterator> StreamIterExt for T { }
44
45pub trait TokenStreamExt
46    : Default
47    + Extend<TokenTree>
48    + Extend<TokenStream>
49    + IntoIterator<Item = TokenTree>
50    + Sized
51{
52    /// Extend a [`TokenTree`]
53    fn push(&mut self, tt: TokenTree) -> &mut Self {
54        self.extend(once(tt));
55        self
56    }
57
58    /// Extend a [`TokenStream`]
59    fn add(&mut self, stream: TokenStream) -> &mut Self {
60        self.extend(once(stream));
61        self
62    }
63
64    /// Call [`mem::take`](std::mem::take)
65    #[must_use]
66    fn take(&mut self) -> Self {
67        take(self)
68    }
69
70    /// Call [`Group::new`]
71    fn grouped(self, delimiter: Delimiter) -> Group;
72
73    fn grouped_paren(self) -> Group {
74        self.grouped(Delimiter::Parenthesis)
75    }
76
77    fn grouped_brace(self) -> Group {
78        self.grouped(Delimiter::Brace)
79    }
80
81    fn grouped_bracket(self) -> Group {
82        self.grouped(Delimiter::Bracket)
83    }
84
85    fn grouped_none(self) -> Group {
86        self.grouped(Delimiter::None)
87    }
88
89    /// Split [`TokenStream`] with `puncts`
90    ///
91    /// Like `"+-,-+".split_puncts(",")` -> `("+-", "-+")`
92    fn split_puncts(self, puncts: impl AsRef<[u8]>) -> Option<(
93        Self,
94        ParseIter<Self::IntoIter>,
95    )>;
96
97    /// Split all [`TokenStream`] with `puncts`
98    ///
99    /// Like `"+-,-+".split_puncts_all(",")` -> `"+-", "-+"`
100    fn split_puncts_all(self, puncts: impl AsRef<[u8]>) -> Vec<Self>;
101}
102impl TokenStreamExt for TokenStream {
103    fn grouped(self, delimiter: Delimiter) -> Group {
104        Group::new(delimiter, self)
105    }
106
107    fn split_puncts(self, puncts: impl AsRef<[u8]>) -> Option<(
108        Self,
109        ParseIter<Self::IntoIter>,
110    )>
111    {
112        let mut iter = self.parse_iter();
113        Some((iter.split_puncts(puncts)?, iter))
114    }
115
116    fn split_puncts_all(self, puncts: impl AsRef<[u8]>) -> Vec<Self> {
117        self.parse_iter()
118            .split_puncts_all(puncts)
119            .collect()
120    }
121
122}
123
124/// Remake each subtree methods
125pub trait WalkExt
126    : IntoIterator<Item = TokenTree>
127    + FromIterator<TokenTree>
128{
129    /// Remake each subtree
130    ///
131    /// `"(1+2)*3"` -> call `f` on `1`, `+`, `2`, `(f(1) f(+) f(2))`, `*`, `3`
132    #[must_use]
133    fn walk<F>(self, mut f: F) -> Self
134    where F: FnMut(TokenTree) -> TokenTree
135    {
136        fn walk_impl<I, F>(this: I, f: &mut F) -> I
137        where I: IntoIterator<Item = TokenTree> + FromIterator<TokenTree>,
138              F: FnMut(TokenTree) -> TokenTree
139        {
140            this.into_iter()
141                .map(|tt| {
142                    let tt = match tt {
143                        TokenTree::Group(g) => {
144                            g.map(|this| walk_impl(this, f)).tt()
145                        },
146                        _ => tt,
147                    };
148                    f(tt)
149                })
150                .collect()
151        }
152        walk_impl(self, &mut f)
153    }
154
155    /// Remake each subtree, and use [`Group`] pack it
156    ///
157    /// Additional call `f` on grouped tokens
158    ///
159    /// `"(1+2)*3"` -> call
160    /// `f` on `1`, `+`, `2`, `(f(1) f(+) f(2))`, `*`, `3`,
161    /// `((f(1) f(+) f(2)) f(*) f(3))`
162    #[must_use]
163    fn grouped_walk<F>(self, delimiter: Delimiter, mut f: F) -> TokenTree
164    where F: FnMut(TokenTree) -> TokenTree,
165    {
166        let stream = self.walk(&mut f);
167        f(Group::new(delimiter, stream.into_iter().collect()).tt())
168    }
169}
170impl<I: IntoIterator<Item = TokenTree> + FromIterator<TokenTree>> WalkExt for I { }
171
172pub trait TokenTreeExt: Into<TokenTree> + Sized {
173    fn as_ident(&self) -> Option<&Ident>     { None }
174    fn as_punct(&self) -> Option<&Punct>     { None }
175    fn as_group(&self) -> Option<&Group>     { None }
176    fn as_literal(&self) -> Option<&Literal> { None }
177    fn into_ident(self) -> Result<Ident, Self>     { Err(self) }
178    fn into_punct(self) -> Result<Punct, Self>     { Err(self) }
179    fn into_group(self) -> Result<Group, Self>     { Err(self) }
180    fn into_literal(self) -> Result<Literal, Self> { Err(self) }
181
182    fn to_ident(&self) -> Result<&Ident, &Self> {
183        self.as_ident().ok_or(self)
184    }
185
186    fn to_punct(&self) -> Result<&Punct, &Self> {
187        self.as_punct().ok_or(self)
188    }
189
190    fn to_group(&self) -> Result<&Group, &Self> {
191        self.as_group().ok_or(self)
192    }
193
194    fn to_literal(&self) -> Result<&Literal, &Self> {
195        self.as_literal().ok_or(self)
196    }
197
198    fn is_ident(&self) -> bool {
199        self.as_ident().is_some()
200    }
201
202    fn is_punct(&self) -> bool {
203        self.as_punct().is_some()
204    }
205
206    fn is_group(&self) -> bool {
207        self.as_group().is_some()
208    }
209
210    fn is_literal(&self) -> bool {
211        self.as_literal().is_some()
212    }
213
214    #[allow(clippy::return_self_not_must_use)]
215    fn map_ident<F: FnOnce(Ident) -> R, R: Into<Self>>(self, op: F) -> Self {
216        self.into_ident()
217            .map(op)
218            .map_or_else(identity, Into::into)
219    }
220
221    #[allow(clippy::return_self_not_must_use)]
222    fn map_punct<F: FnOnce(Punct) -> R, R: Into<Self>>(self, op: F) -> Self {
223        self.into_punct()
224            .map(op)
225            .map_or_else(identity, Into::into)
226    }
227
228    #[allow(clippy::return_self_not_must_use)]
229    fn map_group<F: FnOnce(Group) -> R, R: Into<Self>>(self, op: F) -> Self {
230        self.into_group()
231            .map(op)
232            .map_or_else(identity, Into::into)
233    }
234
235    #[allow(clippy::return_self_not_must_use)]
236    fn map_literal<F: FnOnce(Literal) -> R, R: Into<Self>>(self, op: F) -> Self {
237        self.into_literal()
238            .map(op)
239            .map_or_else(identity, Into::into)
240    }
241
242    #[allow(clippy::return_self_not_must_use)]
243    fn inspect_ident<F: FnOnce(&Ident)>(self, op: F) -> Self {
244        self.as_ident().map(op);
245        self
246    }
247
248    #[allow(clippy::return_self_not_must_use)]
249    fn inspect_punct<F: FnOnce(&Punct)>(self, op: F) -> Self {
250        self.as_punct().map(op);
251        self
252    }
253
254    #[allow(clippy::return_self_not_must_use)]
255    fn inspect_group<F: FnOnce(&Group)>(self, op: F) -> Self {
256        self.as_group().map(op);
257        self
258    }
259
260    #[allow(clippy::return_self_not_must_use)]
261    fn inspect_literal<F: FnOnce(&Literal)>(self, op: F) -> Self {
262        self.as_literal().map(op);
263        self
264    }
265
266    fn kind(&self) -> TokenKind {
267        if self.is_literal() {
268            TokenKind::Literal
269        } else if self.is_punct() {
270            TokenKind::Punct
271        } else if self.is_group() {
272            TokenKind::Group
273        } else if self.is_ident() {
274            TokenKind::Ident
275        } else {
276            unimplemented!()
277        }
278    }
279
280    /// Ident content equal to `keyword` str
281    ///
282    /// Other return `false` when `self` is not [`Ident`]
283    fn is_keyword(&self, keyword: &str) -> bool {
284        self.as_ident().is_some_and(|i| i.to_string() == keyword)
285    }
286
287    /// Punct char equal to `ch`
288    ///
289    /// Other return `false` when `self` is not [`Punct`]
290    fn is_punch(&self, ch: char) -> bool {
291        self.as_punct().is_some_and(|p| p.as_char() == ch)
292    }
293
294    /// Group delimiter is not [`Delimiter::None`]
295    ///
296    /// Other return `false` when `self` is not [`Group`]
297    fn is_solid_group(&self) -> bool {
298        self.as_group().is_some_and(|g| g.is_solid_group())
299    }
300
301    /// Group delimiter equal to `delimiter`
302    ///
303    /// Other return `false` when `self` is not [`Group`]
304    fn is_delimiter(&self, delimiter: Delimiter) -> bool {
305        self.as_group().is_some_and(|g| g.is_delimiter(delimiter))
306    }
307
308    /// Like [`self.is_delimiter(Delimiter::Parenthesis)`](#method.is_delimiter)
309    fn is_delimiter_paren(&self) -> bool {
310        self.is_delimiter(Delimiter::Parenthesis)
311    }
312
313    /// Like [`self.is_delimiter(Delimiter::Brace)`](#method.is_delimiter)
314    fn is_delimiter_brace(&self) -> bool {
315        self.is_delimiter(Delimiter::Brace)
316    }
317
318    /// Like [`self.is_delimiter(Delimiter::Bracket)`](#method.is_delimiter)
319    fn is_delimiter_bracket(&self) -> bool {
320        self.is_delimiter(Delimiter::Bracket)
321    }
322
323    /// Like [`self.is_delimiter(Delimiter::None)`](#method.is_delimiter)
324    fn is_delimiter_none(&self) -> bool {
325        self.is_delimiter(Delimiter::None)
326    }
327
328    /// Call [`Group::stream`] when [`delimiter`] is [`Delimiter::Parenthesis`]
329    ///
330    /// Other return `false` when `self` is not [`Group`]
331    ///
332    /// [`delimiter`]: Group::delimiter
333    fn to_paren_stream(&self) -> Result<TokenStream, &Self> {
334        self.as_group()
335            .and_then(|g| g.is_delimiter_paren().then(|| g.stream()))
336            .ok_or(self)
337    }
338
339    /// Call [`Group::stream`] when [`delimiter`] is [`Delimiter::Brace`]
340    ///
341    /// Other return `false` when `self` is not [`Group`]
342    ///
343    /// [`delimiter`]: Group::delimiter
344    fn to_brace_stream(&self) -> Result<TokenStream, &Self> {
345        self.as_group()
346            .and_then(|g| g.is_delimiter_brace().then(|| g.stream()))
347            .ok_or(self)
348    }
349
350    /// Call [`Group::stream`] when [`delimiter`] is [`Delimiter::Bracket`]
351    ///
352    /// Other return `false` when `self` is not [`Group`]
353    ///
354    /// [`delimiter`]: Group::delimiter
355    fn to_bracket_stream(&self) -> Result<TokenStream, &Self> {
356        self.as_group()
357            .and_then(|g| g.is_delimiter_bracket().then(|| g.stream()))
358            .ok_or(self)
359    }
360
361    /// Call [`Group::stream`] when [`delimiter`] is [`Delimiter::None`]
362    ///
363    /// Other return `false` when `self` is not [`Group`]
364    ///
365    /// [`delimiter`]: Group::delimiter
366    fn to_none_stream(&self) -> Result<TokenStream, &Self> {
367        self.as_group()
368            .and_then(|g| g.is_delimiter_none().then(|| g.stream()))
369            .ok_or(self)
370    }
371
372    /// Like [`to_paren_stream`](#method.to_paren_stream),
373    /// but using `Self` instead of `&Self`
374    fn into_paren_stream(self) -> Result<TokenStream, Self> {
375        self.to_paren_stream()
376            .ok().ok_or(self)
377    }
378
379    /// Like [`to_brace_stream`](#method.to_brace_stream),
380    /// but using `Self` instead of `&Self`
381    fn into_brace_stream(self) -> Result<TokenStream, Self> {
382        self.to_brace_stream()
383            .ok().ok_or(self)
384    }
385
386    /// Like [`to_bracket_stream`](#method.to_bracket_stream),
387    /// but using `Self` instead of `&Self`
388    fn into_bracket_stream(self) -> Result<TokenStream, Self> {
389        self.to_bracket_stream()
390            .ok().ok_or(self)
391    }
392
393    /// Like [`to_none_stream`](#method.to_none_stream),
394    /// but using `Self` instead of `&Self`
395    fn into_none_stream(self) -> Result<TokenStream, Self> {
396        self.to_none_stream()
397            .ok().ok_or(self)
398    }
399
400    /// Punct spacing is [`Joint`]
401    ///
402    /// Other return `false` when `self` is not [`Punct`]
403    fn is_joint(&self) -> bool {
404        self.as_punct().is_some_and(|p| p.spacing() == Joint)
405    }
406
407    fn as_punct_char(&self) -> Option<char> {
408        self.as_punct().map(|p| p.as_char())
409    }
410
411    /// [`Into`] [`TokenTree`], like [`TokenTree::from(self)`]
412    ///
413    /// [`TokenTree::from(self)`]: TokenTree::from
414    fn tt(self) -> TokenTree {
415        self.into()
416    }
417
418    /// [`TokenStream::from_iter(self.tt())`](TokenStream::from_iter)
419    fn unit_stream(self) -> TokenStream {
420        self.tt().into()
421    }
422}
423impl TokenTreeExt for TokenTree {
424    fn as_ident(&self) -> Option<&Ident> {
425        match self {
426            TokenTree::Ident(i) => Some(i),
427            _ => None,
428        }
429    }
430
431    fn as_punct(&self) -> Option<&Punct> {
432        match self {
433            TokenTree::Punct(i) => Some(i),
434            _ => None,
435        }
436    }
437
438    fn as_group(&self) -> Option<&Group> {
439        match self {
440            TokenTree::Group(i) => Some(i),
441            _ => None,
442        }
443    }
444
445    fn as_literal(&self) -> Option<&Literal> {
446        match self {
447            TokenTree::Literal(i) => Some(i),
448            _ => None,
449        }
450    }
451
452    fn into_ident(self) -> Result<Ident, Self> {
453        match self {
454            TokenTree::Ident(i) => Ok(i),
455            _ => Err(self),
456        }
457    }
458
459    fn into_punct(self) -> Result<Punct, Self> {
460        match self {
461            TokenTree::Punct(i) => Ok(i),
462            _ => Err(self),
463        }
464    }
465
466    fn into_group(self) -> Result<Group, Self> {
467        match self {
468            TokenTree::Group(i) => Ok(i),
469            _ => Err(self),
470        }
471    }
472
473    fn into_literal(self) -> Result<Literal, Self> {
474        match self {
475            TokenTree::Literal(i) => Ok(i),
476            _ => Err(self),
477        }
478    }
479
480    fn kind(&self) -> TokenKind {
481        self.into()
482    }
483}
484macro_rules! impl_token_tree_ext {
485    ($as:ident, $into:ident, $ty:ident) => {
486        impl TokenTreeExt for $ty {
487            fn $as(&self) -> Option<&$ty> {
488                Some(self)
489            }
490            fn $into(self) -> Result<$ty, Self> {
491                Ok(self)
492            }
493            fn kind(&self) -> TokenKind {
494                TokenKind::$ty
495            }
496        }
497    };
498}
499impl_token_tree_ext!(as_ident,   into_ident,   Ident);
500impl_token_tree_ext!(as_punct,   into_punct,   Punct);
501impl_token_tree_ext!(as_literal, into_literal, Literal);
502impl TokenTreeExt for Group {
503    fn as_group(&self) -> Option<&Group> {
504        Some(self)
505    }
506
507    fn into_group(self) -> Result<Group, Self> {
508        Ok(self)
509    }
510
511    fn kind(&self) -> TokenKind {
512        TokenKind::Group
513    }
514
515    fn is_solid_group(&self) -> bool {
516        self.delimiter() != Delimiter::None
517    }
518
519    fn is_delimiter(&self, delimiter: Delimiter) -> bool {
520        self.delimiter() == delimiter
521    }
522}
523
524/// Create unsuffixed [`Literal`]
525pub trait Unsuffixed {
526    fn unsuffixed(self) -> Literal;
527}
528/// Create suffixed [`Literal`]
529pub trait Suffixed {
530    fn suffixed(self) -> Literal;
531}
532macro_rules! impl_unsuffixes {
533    ( $($ty:ty: $unsuffixed:ident $($suffixed:ident)?);+ $(;)? ) => {
534        $(
535            #[doc = concat!(
536                "Call [`Literal::",
537                stringify!($unsuffixed),
538                "`]",
539            )]
540            impl Unsuffixed for $ty {
541                fn unsuffixed(self) -> Literal {
542                    Literal::$unsuffixed(self)
543                }
544            }
545
546            $(
547                #[doc = concat!(
548                    "Call [`Literal::",
549                    stringify!($suffixed),
550                    "`]",
551                )]
552                impl Suffixed for $ty {
553                    fn suffixed(self) -> Literal {
554                        Literal::$suffixed(self)
555                    }
556                }
557            )?
558        )*
559    };
560}
561impl_unsuffixes! {
562    i8:     i8_unsuffixed       i8_suffixed;
563    i16:    i16_unsuffixed      i16_suffixed;
564    i32:    i32_unsuffixed      i32_suffixed;
565    i64:    i64_unsuffixed      i64_suffixed;
566    i128:   i128_unsuffixed     i128_suffixed;
567    u8:     u8_unsuffixed       u8_suffixed;
568    u16:    u16_unsuffixed      u16_suffixed;
569    u32:    u32_unsuffixed      u32_suffixed;
570    u64:    u64_unsuffixed      u64_suffixed;
571    u128:   u128_unsuffixed     u128_suffixed;
572    f32:    f32_unsuffixed      f32_suffixed;
573    f64:    f64_unsuffixed      f64_suffixed;
574    usize:  usize_unsuffixed    usize_suffixed;
575    isize:  isize_unsuffixed    isize_suffixed;
576    char:   character;
577    &str:   string;
578    &[u8]:  byte_string;
579}
580
581pub trait PunctsExt: AsRef<[u8]> {
582    /// Call [`puncts`]
583    fn puncts(&self) -> TokenStream {
584        puncts(self)
585    }
586
587    /// Call [`puncts_spanned`]
588    fn puncts_spanned(&self, span: Span) -> TokenStream {
589        puncts_spanned(self, span)
590    }
591}
592impl<T: AsRef<[u8]> + ?Sized> PunctsExt for T { }
593
594pub trait PunctExt: Sized {
595    fn punct(self, spacing: Spacing) -> Punct;
596
597    /// Like [`.punct(Joint)`](#method.punct)
598    fn joint(self) -> Punct {
599        self.punct(Joint)
600    }
601
602    /// Like [`.punct(Alone)`](#method.punct)
603    fn alone(self) -> Punct {
604        self.punct(Alone)
605    }
606}
607impl PunctExt for char {
608    /// Call [`Punct::new`]
609    fn punct(self, spacing: Spacing) -> Punct {
610        Punct::new(self, spacing)
611    }
612}
613
614pub trait StrExt {
615    fn ident(&self, span: Span) -> Ident;
616}
617impl StrExt for str {
618    /// Call [`Ident::new`]
619    fn ident(&self, span: Span) -> Ident {
620        Ident::new(self, span)
621    }
622}
623
624pub trait GroupExt {
625    #[must_use]
626    fn map<F>(&self, f: F) -> Self
627    where F: FnOnce(TokenStream) -> TokenStream;
628}
629impl GroupExt for Group {
630    fn map<F>(&self, f: F) -> Self
631    where F: FnOnce(TokenStream) -> TokenStream,
632    {
633        f(self.stream())
634            .grouped(self.delimiter())
635            .set_spaned(self.span())
636    }
637}
638
639// FIXME: 在下个兼容范围移除punct
640pub trait SpacingExt {
641    /// `Punct::new(ch, self)`
642    #[doc(hidden)]
643    #[deprecated = "renaming to `punch`"]
644    fn punct(self, ch: char) -> Punct where Self: Sized {
645        self.punch(ch)
646    }
647
648    /// `Punct::new(ch, self)`
649    fn punch(self, ch: char) -> Punct where Self: Sized {
650        #[allow(deprecated)]
651        self.punct(ch)
652    }
653
654    /// `*self == Joint`
655    fn is_joint(&self) -> bool;
656
657    /// `*self == Alone`
658    fn is_alone(&self) -> bool;
659}
660impl SpacingExt for Spacing {
661    fn punch(self, ch: char) -> Punct {
662        Punct::new(ch, self)
663    }
664
665    fn is_joint(&self) -> bool {
666        *self == Joint
667    }
668
669    fn is_alone(&self) -> bool {
670        *self == Alone
671    }
672}
proc_macro_tool/exts.rs

proc_macro_tool/
exts.rs
