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    fn as_punch(&self, ch: char) -> Option<&Punct> {
412        self.as_punct()
413            .filter(|it| it.as_char() == ch)
414    }
415
416    fn as_keyword(&self, keyword: &str) -> Option<&Ident> {
417        self.as_ident()
418            .filter(|id| id.is_keyword(keyword))
419    }
420
421    fn to_punch(&self, ch: char) -> Result<&Punct, &Self> {
422        self.as_punch(ch)
423            .ok_or(self)
424    }
425
426    fn to_keyword(&self, keyword: &str) -> Result<&Ident, &Self> {
427        self.as_keyword(keyword)
428            .ok_or(self)
429    }
430
431    fn into_punch(self, ch: char) -> Result<Punct, Self> {
432        if self.is_punch(ch) {
433            self.into_punct()
434        } else {
435            Err(self)
436        }
437    }
438
439    fn into_keyword(self, keyword: &str) -> Result<Ident, Self> {
440        if self.is_keyword(keyword) {
441            self.into_ident()
442        } else {
443            Err(self)
444        }
445    }
446
447    /// [`Into`] [`TokenTree`], like [`TokenTree::from(self)`]
448    ///
449    /// [`TokenTree::from(self)`]: TokenTree::from
450    fn tt(self) -> TokenTree {
451        self.into()
452    }
453
454    /// [`TokenStream::from_iter(self.tt())`](TokenStream::from_iter)
455    fn unit_stream(self) -> TokenStream {
456        self.tt().into()
457    }
458}
459impl TokenTreeExt for TokenTree {
460    fn as_ident(&self) -> Option<&Ident> {
461        match self {
462            TokenTree::Ident(i) => Some(i),
463            _ => None,
464        }
465    }
466
467    fn as_punct(&self) -> Option<&Punct> {
468        match self {
469            TokenTree::Punct(i) => Some(i),
470            _ => None,
471        }
472    }
473
474    fn as_group(&self) -> Option<&Group> {
475        match self {
476            TokenTree::Group(i) => Some(i),
477            _ => None,
478        }
479    }
480
481    fn as_literal(&self) -> Option<&Literal> {
482        match self {
483            TokenTree::Literal(i) => Some(i),
484            _ => None,
485        }
486    }
487
488    fn into_ident(self) -> Result<Ident, Self> {
489        match self {
490            TokenTree::Ident(i) => Ok(i),
491            _ => Err(self),
492        }
493    }
494
495    fn into_punct(self) -> Result<Punct, Self> {
496        match self {
497            TokenTree::Punct(i) => Ok(i),
498            _ => Err(self),
499        }
500    }
501
502    fn into_group(self) -> Result<Group, Self> {
503        match self {
504            TokenTree::Group(i) => Ok(i),
505            _ => Err(self),
506        }
507    }
508
509    fn into_literal(self) -> Result<Literal, Self> {
510        match self {
511            TokenTree::Literal(i) => Ok(i),
512            _ => Err(self),
513        }
514    }
515
516    fn kind(&self) -> TokenKind {
517        self.into()
518    }
519}
520macro_rules! impl_token_tree_ext {
521    ($as:ident, $into:ident, $ty:ident) => {
522        impl TokenTreeExt for $ty {
523            fn $as(&self) -> Option<&$ty> {
524                Some(self)
525            }
526            fn $into(self) -> Result<$ty, Self> {
527                Ok(self)
528            }
529            fn kind(&self) -> TokenKind {
530                TokenKind::$ty
531            }
532        }
533    };
534}
535impl_token_tree_ext!(as_ident,   into_ident,   Ident);
536impl_token_tree_ext!(as_punct,   into_punct,   Punct);
537impl_token_tree_ext!(as_literal, into_literal, Literal);
538impl TokenTreeExt for Group {
539    fn as_group(&self) -> Option<&Group> {
540        Some(self)
541    }
542
543    fn into_group(self) -> Result<Group, Self> {
544        Ok(self)
545    }
546
547    fn kind(&self) -> TokenKind {
548        TokenKind::Group
549    }
550
551    fn is_solid_group(&self) -> bool {
552        self.delimiter() != Delimiter::None
553    }
554
555    fn is_delimiter(&self, delimiter: Delimiter) -> bool {
556        self.delimiter() == delimiter
557    }
558}
559
560/// Create unsuffixed [`Literal`]
561pub trait Unsuffixed {
562    fn unsuffixed(self) -> Literal;
563}
564/// Create suffixed [`Literal`]
565pub trait Suffixed {
566    fn suffixed(self) -> Literal;
567}
568macro_rules! impl_unsuffixes {
569    ( $($ty:ty: $unsuffixed:ident $($suffixed:ident)?);+ $(;)? ) => {
570        $(
571            #[doc = concat!(
572                "Call [`Literal::",
573                stringify!($unsuffixed),
574                "`]",
575            )]
576            impl Unsuffixed for $ty {
577                fn unsuffixed(self) -> Literal {
578                    Literal::$unsuffixed(self)
579                }
580            }
581
582            $(
583                #[doc = concat!(
584                    "Call [`Literal::",
585                    stringify!($suffixed),
586                    "`]",
587                )]
588                impl Suffixed for $ty {
589                    fn suffixed(self) -> Literal {
590                        Literal::$suffixed(self)
591                    }
592                }
593            )?
594        )*
595    };
596}
597impl_unsuffixes! {
598    i8:     i8_unsuffixed       i8_suffixed;
599    i16:    i16_unsuffixed      i16_suffixed;
600    i32:    i32_unsuffixed      i32_suffixed;
601    i64:    i64_unsuffixed      i64_suffixed;
602    i128:   i128_unsuffixed     i128_suffixed;
603    u8:     u8_unsuffixed       u8_suffixed;
604    u16:    u16_unsuffixed      u16_suffixed;
605    u32:    u32_unsuffixed      u32_suffixed;
606    u64:    u64_unsuffixed      u64_suffixed;
607    u128:   u128_unsuffixed     u128_suffixed;
608    f32:    f32_unsuffixed      f32_suffixed;
609    f64:    f64_unsuffixed      f64_suffixed;
610    usize:  usize_unsuffixed    usize_suffixed;
611    isize:  isize_unsuffixed    isize_suffixed;
612    char:   character;
613    &str:   string;
614    &[u8]:  byte_string;
615}
616
617pub trait PunctsExt: AsRef<[u8]> {
618    /// Call [`puncts`]
619    fn puncts(&self) -> TokenStream {
620        puncts(self)
621    }
622
623    /// Call [`puncts_spanned`]
624    fn puncts_spanned(&self, span: Span) -> TokenStream {
625        puncts_spanned(self, span)
626    }
627}
628impl<T: AsRef<[u8]> + ?Sized> PunctsExt for T { }
629
630pub trait PunctExt: Sized {
631    fn punct(self, spacing: Spacing) -> Punct;
632
633    /// Like [`.punct(Joint)`](#method.punct)
634    fn joint(self) -> Punct {
635        self.punct(Joint)
636    }
637
638    /// Like [`.punct(Alone)`](#method.punct)
639    fn alone(self) -> Punct {
640        self.punct(Alone)
641    }
642}
643impl PunctExt for char {
644    /// Call [`Punct::new`]
645    fn punct(self, spacing: Spacing) -> Punct {
646        Punct::new(self, spacing)
647    }
648}
649
650pub trait StrExt {
651    fn ident(&self, span: Span) -> Ident;
652}
653impl StrExt for str {
654    /// Call [`Ident::new`]
655    fn ident(&self, span: Span) -> Ident {
656        Ident::new(self, span)
657    }
658}
659
660pub trait BoolExt {
661    /// Call [`Ident::new`]
662    fn ident(&self) -> Ident;
663}
664impl BoolExt for bool {
665    fn ident(&self) -> Ident {
666        let name = if *self {
667            "true"
668        } else {
669            "false"
670        };
671        Ident::new(name, Span::call_site())
672    }
673}
674
675pub trait GroupExt {
676    #[must_use]
677    fn map<F>(&self, f: F) -> Self
678    where F: FnOnce(TokenStream) -> TokenStream;
679
680    #[must_use]
681    fn map_tts<F>(&self, f: F) -> Self
682    where F: FnMut(TokenTree) -> TokenTree,
683          Self: Sized,
684    {
685        self.map(|stream| stream.into_iter().map(f).collect())
686    }
687
688    #[must_use]
689    fn flat_map_tts<F, I>(&self, f: F) -> Self
690    where F: FnMut(TokenTree) -> I,
691          I: IntoIterator<Item = TokenTree>,
692          Self: Sized,
693    {
694        self.map(|stream| stream.into_iter().flat_map(f).collect())
695    }
696}
697impl GroupExt for Group {
698    fn map<F>(&self, f: F) -> Self
699    where F: FnOnce(TokenStream) -> TokenStream,
700    {
701        f(self.stream())
702            .grouped(self.delimiter())
703            .set_spaned(self.span())
704    }
705}
706
707// FIXME: 在下个兼容范围移除punct
708pub trait SpacingExt {
709    /// `Punct::new(ch, self)`
710    #[doc(hidden)]
711    #[deprecated = "renaming to `punch`"]
712    fn punct(self, ch: char) -> Punct where Self: Sized {
713        self.punch(ch)
714    }
715
716    /// `Punct::new(ch, self)`
717    fn punch(self, ch: char) -> Punct where Self: Sized {
718        #[allow(deprecated)]
719        self.punct(ch)
720    }
721
722    /// `*self == Joint`
723    fn is_joint(&self) -> bool;
724
725    /// `*self == Alone`
726    fn is_alone(&self) -> bool;
727}
728impl SpacingExt for Spacing {
729    fn punch(self, ch: char) -> Punct {
730        Punct::new(ch, self)
731    }
732
733    fn is_joint(&self) -> bool {
734        *self == Joint
735    }
736
737    fn is_alone(&self) -> bool {
738        *self == Alone
739    }
740}
proc_macro_tool/exts.rs

proc_macro_tool/
exts.rs
