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}
155impl<I: IntoIterator<Item = TokenTree> + FromIterator<TokenTree>> WalkExt for I { }
156
157pub trait TokenTreeExt: Into<TokenTree> + Sized {
158    fn as_ident(&self) -> Option<&Ident>     { None }
159    fn as_punct(&self) -> Option<&Punct>     { None }
160    fn as_group(&self) -> Option<&Group>     { None }
161    fn as_literal(&self) -> Option<&Literal> { None }
162    fn into_ident(self) -> Result<Ident, Self>     { Err(self) }
163    fn into_punct(self) -> Result<Punct, Self>     { Err(self) }
164    fn into_group(self) -> Result<Group, Self>     { Err(self) }
165    fn into_literal(self) -> Result<Literal, Self> { Err(self) }
166
167    fn to_ident(&self) -> Result<&Ident, &Self> {
168        self.as_ident().ok_or(self)
169    }
170
171    fn to_punct(&self) -> Result<&Punct, &Self> {
172        self.as_punct().ok_or(self)
173    }
174
175    fn to_group(&self) -> Result<&Group, &Self> {
176        self.as_group().ok_or(self)
177    }
178
179    fn to_literal(&self) -> Result<&Literal, &Self> {
180        self.as_literal().ok_or(self)
181    }
182
183    fn is_ident(&self) -> bool {
184        self.as_ident().is_some()
185    }
186
187    fn is_punct(&self) -> bool {
188        self.as_punct().is_some()
189    }
190
191    fn is_group(&self) -> bool {
192        self.as_group().is_some()
193    }
194
195    fn is_literal(&self) -> bool {
196        self.as_literal().is_some()
197    }
198
199    #[allow(clippy::return_self_not_must_use)]
200    fn map_ident<F: FnOnce(Ident) -> R, R: Into<Self>>(self, op: F) -> Self {
201        self.into_ident()
202            .map(op)
203            .map_or_else(identity, Into::into)
204    }
205
206    #[allow(clippy::return_self_not_must_use)]
207    fn map_punct<F: FnOnce(Punct) -> R, R: Into<Self>>(self, op: F) -> Self {
208        self.into_punct()
209            .map(op)
210            .map_or_else(identity, Into::into)
211    }
212
213    #[allow(clippy::return_self_not_must_use)]
214    fn map_group<F: FnOnce(Group) -> R, R: Into<Self>>(self, op: F) -> Self {
215        self.into_group()
216            .map(op)
217            .map_or_else(identity, Into::into)
218    }
219
220    #[allow(clippy::return_self_not_must_use)]
221    fn map_literal<F: FnOnce(Literal) -> R, R: Into<Self>>(self, op: F) -> Self {
222        self.into_literal()
223            .map(op)
224            .map_or_else(identity, Into::into)
225    }
226
227    fn kind(&self) -> TokenKind {
228        if self.is_literal() {
229            TokenKind::Literal
230        } else if self.is_punct() {
231            TokenKind::Punct
232        } else if self.is_group() {
233            TokenKind::Group
234        } else if self.is_ident() {
235            TokenKind::Ident
236        } else {
237            unimplemented!()
238        }
239    }
240
241    /// Ident content equal to `keyword` str
242    ///
243    /// Other return `false` when `self` is not [`Ident`]
244    fn is_keyword(&self, keyword: &str) -> bool {
245        self.as_ident().is_some_and(|i| i.to_string() == keyword)
246    }
247
248    /// Punct char equal to `ch`
249    ///
250    /// Other return `false` when `self` is not [`Punct`]
251    fn is_punch(&self, ch: char) -> bool {
252        self.as_punct().is_some_and(|p| p.as_char() == ch)
253    }
254
255    /// Group delimiter is not [`Delimiter::None`]
256    ///
257    /// Other return `false` when `self` is not [`Group`]
258    fn is_solid_group(&self) -> bool {
259        self.as_group().is_some_and(|g| g.is_solid_group())
260    }
261
262    /// Group delimiter equal to `delimiter`
263    ///
264    /// Other return `false` when `self` is not [`Group`]
265    fn is_delimiter(&self, delimiter: Delimiter) -> bool {
266        self.as_group().is_some_and(|g| g.is_delimiter(delimiter))
267    }
268
269    /// Like [`self.is_delimiter(Delimiter::Parenthesis)`](#method.is_delimiter)
270    fn is_delimiter_paren(&self) -> bool {
271        self.is_delimiter(Delimiter::Parenthesis)
272    }
273
274    /// Like [`self.is_delimiter(Delimiter::Brace)`](#method.is_delimiter)
275    fn is_delimiter_brace(&self) -> bool {
276        self.is_delimiter(Delimiter::Brace)
277    }
278
279    /// Like [`self.is_delimiter(Delimiter::Bracket)`](#method.is_delimiter)
280    fn is_delimiter_bracket(&self) -> bool {
281        self.is_delimiter(Delimiter::Bracket)
282    }
283
284    /// Like [`self.is_delimiter(Delimiter::None)`](#method.is_delimiter)
285    fn is_delimiter_none(&self) -> bool {
286        self.is_delimiter(Delimiter::None)
287    }
288
289    /// Call [`Group::stream`] when [`delimiter`] is [`Delimiter::Parenthesis`]
290    ///
291    /// Other return `false` when `self` is not [`Group`]
292    ///
293    /// [`delimiter`]: Group::delimiter
294    fn to_paren_stream(&self) -> Result<TokenStream, &Self> {
295        self.as_group()
296            .and_then(|g| g.is_delimiter_paren().then(|| g.stream()))
297            .ok_or(self)
298    }
299
300    /// Call [`Group::stream`] when [`delimiter`] is [`Delimiter::Brace`]
301    ///
302    /// Other return `false` when `self` is not [`Group`]
303    ///
304    /// [`delimiter`]: Group::delimiter
305    fn to_brace_stream(&self) -> Result<TokenStream, &Self> {
306        self.as_group()
307            .and_then(|g| g.is_delimiter_brace().then(|| g.stream()))
308            .ok_or(self)
309    }
310
311    /// Call [`Group::stream`] when [`delimiter`] is [`Delimiter::Bracket`]
312    ///
313    /// Other return `false` when `self` is not [`Group`]
314    ///
315    /// [`delimiter`]: Group::delimiter
316    fn to_bracket_stream(&self) -> Result<TokenStream, &Self> {
317        self.as_group()
318            .and_then(|g| g.is_delimiter_bracket().then(|| g.stream()))
319            .ok_or(self)
320    }
321
322    /// Call [`Group::stream`] when [`delimiter`] is [`Delimiter::None`]
323    ///
324    /// Other return `false` when `self` is not [`Group`]
325    ///
326    /// [`delimiter`]: Group::delimiter
327    fn to_none_stream(&self) -> Result<TokenStream, &Self> {
328        self.as_group()
329            .and_then(|g| g.is_delimiter_none().then(|| g.stream()))
330            .ok_or(self)
331    }
332
333    /// Like [`to_paren_stream`](#method.to_paren_stream),
334    /// but using `Self` instead of `&Self`
335    fn into_paren_stream(self) -> Result<TokenStream, Self> {
336        self.to_paren_stream()
337            .ok().ok_or(self)
338    }
339
340    /// Like [`to_brace_stream`](#method.to_brace_stream),
341    /// but using `Self` instead of `&Self`
342    fn into_brace_stream(self) -> Result<TokenStream, Self> {
343        self.to_brace_stream()
344            .ok().ok_or(self)
345    }
346
347    /// Like [`to_bracket_stream`](#method.to_bracket_stream),
348    /// but using `Self` instead of `&Self`
349    fn into_bracket_stream(self) -> Result<TokenStream, Self> {
350        self.to_bracket_stream()
351            .ok().ok_or(self)
352    }
353
354    /// Like [`to_none_stream`](#method.to_none_stream),
355    /// but using `Self` instead of `&Self`
356    fn into_none_stream(self) -> Result<TokenStream, Self> {
357        self.to_none_stream()
358            .ok().ok_or(self)
359    }
360
361    /// Punct spacing is [`Joint`]
362    ///
363    /// Other return `false` when `self` is not [`Punct`]
364    fn is_joint(&self) -> bool {
365        self.as_punct().is_some_and(|p| p.spacing() == Joint)
366    }
367
368    fn as_punct_char(&self) -> Option<char> {
369        self.as_punct().map(|p| p.as_char())
370    }
371
372    /// [`Into`] [`TokenTree`], like [`TokenTree::from(self)`]
373    ///
374    /// [`TokenTree::from(self)`]: TokenTree::from
375    fn tt(self) -> TokenTree {
376        self.into()
377    }
378
379    /// [`TokenStream::from_iter(self.tt())`](TokenStream::from_iter)
380    fn unit_stream(self) -> TokenStream {
381        self.tt().into()
382    }
383}
384impl TokenTreeExt for TokenTree {
385    fn as_ident(&self) -> Option<&Ident> {
386        match self {
387            TokenTree::Ident(i) => Some(i),
388            _ => None,
389        }
390    }
391
392    fn as_punct(&self) -> Option<&Punct> {
393        match self {
394            TokenTree::Punct(i) => Some(i),
395            _ => None,
396        }
397    }
398
399    fn as_group(&self) -> Option<&Group> {
400        match self {
401            TokenTree::Group(i) => Some(i),
402            _ => None,
403        }
404    }
405
406    fn as_literal(&self) -> Option<&Literal> {
407        match self {
408            TokenTree::Literal(i) => Some(i),
409            _ => None,
410        }
411    }
412
413    fn into_ident(self) -> Result<Ident, Self> {
414        match self {
415            TokenTree::Ident(i) => Ok(i),
416            _ => Err(self),
417        }
418    }
419
420    fn into_punct(self) -> Result<Punct, Self> {
421        match self {
422            TokenTree::Punct(i) => Ok(i),
423            _ => Err(self),
424        }
425    }
426
427    fn into_group(self) -> Result<Group, Self> {
428        match self {
429            TokenTree::Group(i) => Ok(i),
430            _ => Err(self),
431        }
432    }
433
434    fn into_literal(self) -> Result<Literal, Self> {
435        match self {
436            TokenTree::Literal(i) => Ok(i),
437            _ => Err(self),
438        }
439    }
440
441    fn kind(&self) -> TokenKind {
442        self.into()
443    }
444}
445macro_rules! impl_token_tree_ext {
446    ($as:ident, $into:ident, $ty:ident) => {
447        impl TokenTreeExt for $ty {
448            fn $as(&self) -> Option<&$ty> {
449                Some(self)
450            }
451            fn $into(self) -> Result<$ty, Self> {
452                Ok(self)
453            }
454            fn kind(&self) -> TokenKind {
455                TokenKind::$ty
456            }
457        }
458    };
459}
460impl_token_tree_ext!(as_ident,   into_ident,   Ident);
461impl_token_tree_ext!(as_punct,   into_punct,   Punct);
462impl_token_tree_ext!(as_literal, into_literal, Literal);
463impl TokenTreeExt for Group {
464    fn as_group(&self) -> Option<&Group> {
465        Some(self)
466    }
467
468    fn into_group(self) -> Result<Group, Self> {
469        Ok(self)
470    }
471
472    fn kind(&self) -> TokenKind {
473        TokenKind::Group
474    }
475
476    fn is_solid_group(&self) -> bool {
477        self.delimiter() != Delimiter::None
478    }
479
480    fn is_delimiter(&self, delimiter: Delimiter) -> bool {
481        self.delimiter() == delimiter
482    }
483}
484
485/// Create unsuffixed [`Literal`]
486pub trait Unsuffixed {
487    fn unsuffixed(self) -> Literal;
488}
489/// Create suffixed [`Literal`]
490pub trait Suffixed {
491    fn suffixed(self) -> Literal;
492}
493macro_rules! impl_unsuffixes {
494    ( $($ty:ty: $unsuffixed:ident $($suffixed:ident)?);+ $(;)? ) => {
495        $(
496            #[doc = concat!(
497                "Call [`Literal::",
498                stringify!($unsuffixed),
499                "`]",
500            )]
501            impl Unsuffixed for $ty {
502                fn unsuffixed(self) -> Literal {
503                    Literal::$unsuffixed(self)
504                }
505            }
506
507            $(
508                #[doc = concat!(
509                    "Call [`Literal::",
510                    stringify!($suffixed),
511                    "`]",
512                )]
513                impl Suffixed for $ty {
514                    fn suffixed(self) -> Literal {
515                        Literal::$suffixed(self)
516                    }
517                }
518            )?
519        )*
520    };
521}
522impl_unsuffixes! {
523    i8:     i8_unsuffixed       i8_suffixed;
524    i16:    i16_unsuffixed      i16_suffixed;
525    i32:    i32_unsuffixed      i32_suffixed;
526    i64:    i64_unsuffixed      i64_suffixed;
527    i128:   i128_unsuffixed     i128_suffixed;
528    u8:     u8_unsuffixed       u8_suffixed;
529    u16:    u16_unsuffixed      u16_suffixed;
530    u32:    u32_unsuffixed      u32_suffixed;
531    u64:    u64_unsuffixed      u64_suffixed;
532    u128:   u128_unsuffixed     u128_suffixed;
533    f32:    f32_unsuffixed      f32_suffixed;
534    f64:    f64_unsuffixed      f64_suffixed;
535    usize:  usize_unsuffixed    usize_suffixed;
536    isize:  isize_unsuffixed    isize_suffixed;
537    char:   character;
538    &str:   string;
539    &[u8]:  byte_string;
540}
541
542pub trait PunctsExt: AsRef<[u8]> {
543    /// Call [`puncts`]
544    fn puncts(&self) -> TokenStream {
545        puncts(self)
546    }
547
548    /// Call [`puncts_spanned`]
549    fn puncts_spanned(&self, span: Span) -> TokenStream {
550        puncts_spanned(self, span)
551    }
552}
553impl<T: AsRef<[u8]> + ?Sized> PunctsExt for T { }
554
555pub trait PunctExt: Sized {
556    fn punct(self, spacing: Spacing) -> Punct;
557
558    /// Like [`.punct(Joint)`](#method.punct)
559    fn joint(self) -> Punct {
560        self.punct(Joint)
561    }
562
563    /// Like [`.punct(Alone)`](#method.punct)
564    fn alone(self) -> Punct {
565        self.punct(Alone)
566    }
567}
568impl PunctExt for char {
569    /// Call [`Punct::new`]
570    fn punct(self, spacing: Spacing) -> Punct {
571        Punct::new(self, spacing)
572    }
573}
574
575pub trait StrExt {
576    fn ident(&self, span: Span) -> Ident;
577}
578impl StrExt for str {
579    /// Call [`Ident::new`]
580    fn ident(&self, span: Span) -> Ident {
581        Ident::new(self, span)
582    }
583}
584
585pub trait GroupExt {
586    #[must_use]
587    fn map<F>(&self, f: F) -> Self
588    where F: FnOnce(TokenStream) -> TokenStream;
589}
590impl GroupExt for Group {
591    fn map<F>(&self, f: F) -> Self
592    where F: FnOnce(TokenStream) -> TokenStream,
593    {
594        f(self.stream())
595            .grouped(self.delimiter())
596            .set_spaned(self.span())
597    }
598}
599
600// FIXME: 在下个兼容范围移除punct
601pub trait SpacingExt {
602    /// `Punct::new(ch, self)`
603    #[doc(hidden)]
604    #[deprecated = "renaming to `punch`"]
605    fn punct(self, ch: char) -> Punct where Self: Sized {
606        self.punch(ch)
607    }
608
609    /// `Punct::new(ch, self)`
610    fn punch(self, ch: char) -> Punct where Self: Sized {
611        #[allow(deprecated)]
612        self.punct(ch)
613    }
614
615    /// `*self == Joint`
616    fn is_joint(&self) -> bool;
617
618    /// `*self == Alone`
619    fn is_alone(&self) -> bool;
620}
621impl SpacingExt for Spacing {
622    fn punch(self, ch: char) -> Punct {
623        Punct::new(ch, self)
624    }
625
626    fn is_joint(&self) -> bool {
627        *self == Joint
628    }
629
630    fn is_alone(&self) -> bool {
631        *self == Alone
632    }
633}
proc_macro_tool/exts.rs

proc_macro_tool/
exts.rs
