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

proc_macro_tool/
exts.rs
