proc_macro_tool/
lib.rs

1#![doc = include_str!("../README.md")]
2
3extern crate proc_macro;
4
5use std::{
6    collections::VecDeque,
7    iter::{once, FusedIterator},
8};
9
10use proc_macro::{
11    Delimiter, Group, Ident, Literal, Punct,
12    Spacing::*,
13    Span, TokenStream, TokenTree,
14};
15
16/// Generate a function, set input `TokenTree` span
17pub fn span_setter<T>(span: Span) -> impl Fn(T) -> T
18where T: SetSpan,
19{
20    move |tt| {
21        tt.set_spaned(span)
22    }
23}
24
25pub trait SetSpan: Sized {
26    /// Call [`TokenTree::set_span`]
27    fn set_span(&mut self, span: Span);
28
29    fn set_spaned(mut self, span: Span) -> Self {
30        self.set_span(span);
31        self
32    }
33}
34macro_rules! impl_set_span {
35    ($ty:ty) => {
36        impl SetSpan for $ty {
37            fn set_span(&mut self, span: Span) {
38                self.set_span(span);
39            }
40        }
41    };
42}
43impl_set_span!(TokenTree);
44impl_set_span!(Ident);
45impl_set_span!(Punct);
46impl_set_span!(Group);
47impl_set_span!(Literal);
48
49/// `<TokenStream as FromIterator<TokenTree>>::from_iter`
50#[must_use]
51pub fn stream<I>(iter: I) -> TokenStream
52where I: IntoIterator<Item = TokenTree>,
53{
54    TokenStream::from_iter(iter)
55}
56
57/// `<TokenStream as FromIterator<TokenStream>>::from_iter`
58#[must_use]
59pub fn streams<I>(iter: I) -> TokenStream
60where I: IntoIterator<Item = TokenStream>,
61{
62    TokenStream::from_iter(iter)
63}
64
65/// Make `compile_error! {"..."}`
66#[must_use]
67pub fn err(msg: &str, span: Span) -> TokenStream {
68    let s = span_setter(span);
69    stream([
70        s(Punct::new(':', Joint).into()),
71        s(Punct::new(':', Joint).into()),
72        s(Ident::new("core", span).into()),
73        s(Punct::new(':', Joint).into()),
74        s(Punct::new(':', Joint).into()),
75        s(Ident::new("compile_error", span).into()),
76        s(Punct::new('!', Joint).into()),
77        s(Group::new(Delimiter::Brace, stream([
78            s(Literal::string(msg).into()),
79        ])).into()),
80    ])
81}
82
83/// Like [`err()`], but use [`Result`]
84///
85/// # Errors
86/// - always return [`Err`]
87pub fn rerr<T>(msg: &str, span: Span) -> Result<T, TokenStream> {
88    Err(err(msg, span))
89}
90
91/// Make puncts, `spacing` is last punct spacing
92///
93/// - `"+-"` like `[Joint('+'), Joint('-')]`
94/// - `"+- "` like `[Joint('+'), Alone('-')]`
95/// - `"+ -"` like `[Alone('+'), Joint('-')]`
96pub fn puncts(puncts: impl AsRef<[u8]>) -> TokenStream {
97    puncts_spanned(puncts, Span::call_site())
98}
99
100/// Make puncts, `spacing` is last punct spacing
101///
102/// Like [`puncts`], but `.set_span(span)`
103pub fn puncts_spanned(puncts: impl AsRef<[u8]>, span: Span) -> TokenStream {
104    let puncts = puncts.as_ref().trim_ascii_start();
105    let iter = &mut puncts.iter().copied().peekable();
106    let mut result = TokenStream::new();
107
108    while let Some(ch) = iter.next() {
109        debug_assert!(! ch.is_ascii_whitespace());
110        let mut s = None;
111        while iter.next_if(u8::is_ascii_whitespace).is_some() {
112            s = Some(Alone)
113        }
114        let spacing = s.or(iter.peek().map(|_| Joint))
115            .unwrap_or(Joint);
116        let p = Punct::new(ch.into(), spacing);
117        result.push(p.set_spaned(span).into());
118    }
119
120    result
121}
122
123/// [`return err(msg [, span])`](err())
124#[macro_export]
125macro_rules! err {
126    ($msg:expr $(,)?) => { $crate::err!($msg, ::proc_macro::Span::call_site()) };
127    ($msg:expr , $span:expr $(,)?) => {
128        return $crate::err($msg, $span)
129    };
130}
131
132/// [`return rerr(msg [, span])`](rerr())
133#[macro_export]
134macro_rules! rerr {
135    ($msg:expr $(,)?) => { $crate::rerr!($msg, ::proc_macro::Span::call_site()) };
136    ($msg:expr , $span:expr $(,)?) => {
137        return $crate::rerr($msg, $span)
138    };
139}
140
141pub trait TokenStreamExt
142    : Default
143    + Extend<TokenTree>
144    + Extend<TokenStream>
145    + IntoIterator<Item = TokenTree>
146    + FromIterator<TokenTree>
147    + Sized
148{
149    fn push(&mut self, tt: TokenTree) -> &mut Self {
150        self.extend(once(tt));
151        self
152    }
153
154    fn add(&mut self, stream: TokenStream) -> &mut Self {
155        self.extend(once(stream));
156        self
157    }
158
159    fn grouped(self, delimiter: Delimiter) -> Group;
160
161    fn walk<F>(self, mut f: F) -> Self
162    where F: FnMut(TokenTree) -> TokenTree
163    {
164        fn walk_impl<T, F>(this: T, f: &mut F) -> T
165        where T: TokenStreamExt,
166              F: FnMut(TokenTree) -> TokenTree
167        {
168            this.into_iter()
169                .map(|tt| {
170                    let tt = match tt {
171                        TokenTree::Group(g) => {
172                            walk_impl(g.stream(), &mut *f)
173                                .grouped(g.delimiter())
174                                .set_spaned(g.span())
175                                .into()
176                        },
177                        _ => tt,
178                    };
179                    f(tt)
180                })
181                .collect()
182        }
183        walk_impl(self, &mut f)
184    }
185
186    /// Split [`TokenStream`] to `predicate` false and true
187    ///
188    /// Like `"+-,-+".split_puncts(",")` -> `("+-", "-+")`
189    fn split_puncts(self, puncts: impl AsRef<[u8]>) -> Option<(
190        Self,
191        ParseIter<Self::IntoIter>,
192    )>;
193}
194impl TokenStreamExt for TokenStream {
195    fn grouped(self, delimiter: Delimiter) -> Group {
196        Group::new(delimiter, self)
197    }
198
199    fn split_puncts(self, puncts: impl AsRef<[u8]>) -> Option<(
200        Self,
201        ParseIter<Self::IntoIter>,
202    )>
203    {
204        let mut iter = self.parse_iter();
205        Some((iter.split_puncts(puncts)?, iter))
206    }
207
208}
209
210pub trait WalkExt
211    : IntoIterator<Item = TokenTree>
212    + FromIterator<TokenTree>
213{
214    /// Remake each subtree
215    ///
216    /// `"(1+2)*3"` -> call `f` on `1`, `+`, `2`, `(1+2)`, `*`, `3`
217    fn walk<F>(self, mut f: F) -> Self
218    where F: FnMut(TokenTree) -> TokenTree
219    {
220        fn walk_impl<I, F>(this: I, f: &mut F) -> I
221        where I: IntoIterator<Item = TokenTree> + FromIterator<TokenTree>,
222              F: FnMut(TokenTree) -> TokenTree
223        {
224            this.into_iter()
225                .map(|tt| {
226                    let tt = match tt {
227                        TokenTree::Group(g) => {
228                            walk_impl(g.stream(), &mut *f)
229                                .grouped(g.delimiter())
230                                .set_spaned(g.span())
231                                .into()
232                        },
233                        _ => tt,
234                    };
235                    f(tt)
236                })
237                .collect()
238        }
239        walk_impl(self, &mut f)
240    }
241}
242
243pub trait TokenTreeExt: Sized {
244    fn as_ident(&self) -> Option<&Ident>;
245    fn as_punct(&self) -> Option<&Punct>;
246    fn as_group(&self) -> Option<&Group>;
247    fn as_literal(&self) -> Option<&Literal>;
248    fn into_ident(self) -> Result<Ident, Self>;
249    fn into_punct(self) -> Result<Punct, Self>;
250    fn into_group(self) -> Result<Group, Self>;
251    fn into_literal(self) -> Result<Literal, Self>;
252
253    fn is_ident(&self) -> bool {
254        self.as_ident().is_some()
255    }
256
257    fn is_punct(&self) -> bool {
258        self.as_punct().is_some()
259    }
260
261    fn is_group(&self) -> bool {
262        self.as_group().is_some()
263    }
264
265    fn is_literal(&self) -> bool {
266        self.as_literal().is_some()
267    }
268
269    /// Ident content equal to `keyword` str
270    ///
271    /// Other return `false` when `self` is not [`Ident`]
272    fn is_keyword(&self, keyword: &str) -> bool {
273        self.as_ident().is_some_and(|i| i.to_string() == keyword)
274    }
275
276    /// Punct char equal to `ch`
277    ///
278    /// Other return `false` when `self` is not [`Punct`]
279    fn is_punch(&self, ch: char) -> bool {
280        self.as_punct().is_some_and(|p| p.as_char() == ch)
281    }
282
283    /// Group delimiter is not [`Delimiter::None`]
284    ///
285    /// Other return `false` when `self` is not [`Group`]
286    fn is_solid_group(&self) -> bool {
287        self.as_group().is_some_and(|g| g.is_solid_group())
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}
301impl TokenTreeExt for TokenTree {
302    fn as_ident(&self) -> Option<&Ident> {
303        match self {
304            TokenTree::Ident(i) => Some(i),
305            _ => None,
306        }
307    }
308
309    fn as_punct(&self) -> Option<&Punct> {
310        match self {
311            TokenTree::Punct(i) => Some(i),
312            _ => None,
313        }
314    }
315
316    fn as_group(&self) -> Option<&Group> {
317        match self {
318            TokenTree::Group(i) => Some(i),
319            _ => None,
320        }
321    }
322
323    fn as_literal(&self) -> Option<&Literal> {
324        match self {
325            TokenTree::Literal(i) => Some(i),
326            _ => None,
327        }
328    }
329
330    fn into_ident(self) -> Result<Ident, Self> {
331        match self {
332            TokenTree::Ident(i) => Ok(i),
333            _ => Err(self),
334        }
335    }
336
337    fn into_punct(self) -> Result<Punct, Self> {
338        match self {
339            TokenTree::Punct(i) => Ok(i),
340            _ => Err(self),
341        }
342    }
343
344    fn into_group(self) -> Result<Group, Self> {
345        match self {
346            TokenTree::Group(i) => Ok(i),
347            _ => Err(self),
348        }
349    }
350
351    fn into_literal(self) -> Result<Literal, Self> {
352        match self {
353            TokenTree::Literal(i) => Ok(i),
354            _ => Err(self),
355        }
356    }
357}
358
359pub trait GroupExt {
360    /// Group delimiter is not [`Delimiter::None`]
361    fn is_solid_group(&self) -> bool;
362}
363impl GroupExt for Group {
364    fn is_solid_group(&self) -> bool {
365        self.delimiter() != Delimiter::None
366    }
367}
368
369#[derive(Debug, Clone)]
370pub struct ParseIter<I: Iterator<Item = TokenTree>> {
371    iter: I,
372    buf: VecDeque<TokenTree>,
373}
374
375impl<I: Iterator<Item = TokenTree>> ParseIter<I> {
376    pub fn peek(&mut self) -> Option<&TokenTree> {
377        self.peek_i(0)
378    }
379
380    pub fn next_if<F>(&mut self, f: F) -> Option<TokenTree>
381    where F: FnOnce(&TokenTree) -> bool,
382    {
383        let peek = self.peek()?;
384
385        if f(peek) {
386            self.next()
387        } else {
388            None
389        }
390    }
391
392    pub fn peek_i(&mut self, i: usize) -> Option<&TokenTree> {
393        for _ in self.buf.len()..=i {
394            self.buf.push_back(self.iter.next()?);
395        }
396        Some(&self.buf[i])
397    }
398
399    pub fn peek_is<F>(&mut self, f: F) -> bool
400    where F: FnOnce(&TokenTree) -> bool,
401    {
402        self.peek().is_some_and(f)
403    }
404
405    pub fn peek_i_is<F>(&mut self, i: usize, f: F) -> bool
406    where F: FnOnce(&TokenTree) -> bool,
407    {
408        self.peek_i(i).is_some_and(f)
409    }
410
411    /// Peek jointed puncts
412    pub fn peek_puncts(&mut self, puncts: impl AsRef<[u8]>) -> Option<
413        impl Iterator<Item = &TokenTree>
414    > {
415        let mut prev = None;
416
417        for (i, ch) in puncts.as_ref().iter()
418            .copied().map(char::from).enumerate()
419        {
420            if let Some(prev) = prev {
421                if prev == Alone { return None }
422            }
423
424            let tt = self.peek_i(i)?;
425            let TokenTree::Punct(p) = tt else { return None };
426            if p.as_char() != ch { return None }
427
428            prev = Some(p.spacing());
429        }
430        Some(self.buf.iter())
431    }
432
433    /// Next jointed puncts
434    pub fn next_puncts(&mut self, puncts: impl AsRef<[u8]>) -> Option<
435        impl Iterator<Item = TokenTree> + '_
436    > {
437        let _ = self.peek_puncts(puncts.as_ref())?;
438        Some(self.buf.drain(..puncts.as_ref().len()))
439    }
440
441    /// Split [`TokenStream`] to `predicate` false and true
442    ///
443    /// Like `"+-,-+".split_puncts(",")` -> `("+-", "-+")`
444    pub fn split_puncts(&mut self, puncts: impl AsRef<[u8]>) -> Option<TokenStream> {
445        let mut left = TokenStream::new();
446        let puncts = puncts.as_ref();
447
448        loop {
449            if self.next_puncts(puncts).is_some() {
450                break Some(left);
451            }
452            if self.peek().is_none() {
453                break None;
454            }
455            left.push(self.next().unwrap());
456        }
457    }
458}
459pub trait ParseIterExt: IntoIterator<Item = TokenTree> + Sized {
460    fn parse_iter(self) -> ParseIter<Self::IntoIter> {
461        ParseIter { iter: self.into_iter(), buf: VecDeque::new() }
462    }
463}
464impl<I: IntoIterator<Item = TokenTree>> ParseIterExt for I { }
465
466impl<I: Iterator<Item = TokenTree>> Iterator for ParseIter<I> {
467    type Item = TokenTree;
468
469    fn next(&mut self) -> Option<Self::Item> {
470        self.buf.pop_front()
471            .or_else(|| self.iter.next())
472    }
473
474    fn count(self) -> usize
475    where Self: Sized,
476    {
477        self.buf.len() + self.iter.count()
478    }
479
480    fn fold<B, F>(self, init: B, f: F) -> B
481    where Self: Sized,
482          F: FnMut(B, Self::Item) -> B,
483    {
484        self.buf.into_iter().chain(self.iter).fold(init, f)
485    }
486
487    fn size_hint(&self) -> (usize, Option<usize>) {
488        let (lo, hi) = self.iter.size_hint();
489        let lo = lo.saturating_add(self.buf.len());
490        let hi = hi.and_then(|hi| hi.checked_add(self.buf.len()));
491        (lo, hi)
492    }
493}
494impl<I: DoubleEndedIterator<Item = TokenTree>> DoubleEndedIterator for ParseIter<I> {
495    fn next_back(&mut self) -> Option<Self::Item> {
496        self.iter.next_back()
497            .or_else(|| self.buf.pop_back())
498    }
499
500    fn rfold<B, F>(self, init: B, f: F) -> B
501    where Self: Sized,
502          F: FnMut(B, Self::Item) -> B,
503    {
504        self.buf.into_iter().chain(self.iter).rfold(init, f)
505    }
506}
507impl<I: ExactSizeIterator<Item = TokenTree>> ExactSizeIterator for ParseIter<I> { }
508impl<I: FusedIterator<Item = TokenTree>> FusedIterator for ParseIter<I> { }
509
510fn pfunc_impl<F, R>(
511    stream: TokenStream,
512    proc_input: bool,
513    names: &[&str],
514    f: &mut F,
515) -> Result<TokenStream, R>
516where F: FnMut(Ident, Group) -> Result<TokenStream, R>,
517{
518    let mut iter = stream.into_iter().parse_iter();
519    let mut result = TokenStream::new();
520
521    while let Some(tt) = iter.next() {
522        match tt {
523            TokenTree::Punct(p)
524                if p.as_char() == '#'
525                && iter.peek_is(|i| i.as_ident()
526                    .is_some_and(|i| names.contains(&&*i.to_string())))
527                && iter.peek_i_is(1, |t| t.is_solid_group())
528                =>
529            {
530                let ident = iter.next().unwrap().into_ident().unwrap();
531                let mut group = iter.next().unwrap().into_group().unwrap();
532                if proc_input {
533                    let sub = pfunc_impl(
534                        group.stream(),
535                        proc_input,
536                        names,
537                        f,
538                    )?;
539                    group = sub
540                        .grouped(group.delimiter())
541                        .set_spaned(group.span());
542                }
543                result.add(f(ident, group)?);
544            },
545            TokenTree::Group(g) => {
546                let sub = pfunc_impl(g.stream(), proc_input, names, f)?;
547                result.push(sub
548                    .grouped(g.delimiter())
549                    .set_spaned(g.span())
550                    .into());
551            },
552            _ => _ = result.push(tt),
553        }
554    }
555
556    Ok(result)
557}
558
559/// Call `f` on `#name(...)` `#name[...]` etc, exclude [`Delimiter::None`]
560///
561/// Apply pfunc for `(...)` when `proc_input` is `true`
562pub fn pfunc<'a>(
563    stream: TokenStream,
564    proc_input: bool,
565    names: impl AsRef<[&'a str]>,
566    mut f: impl FnMut(Ident, Group) -> TokenStream,
567) -> TokenStream {
568    let f = &mut |i, g| {
569        Ok::<_, ()>(f(i, g))
570    };
571    pfunc_impl(stream, proc_input, names.as_ref(), f).unwrap()
572}
573
574/// Call `f` on `#name(...)` `#name[...]` etc, exclude [`Delimiter::None`]
575///
576/// Apply pfunc for `(...)` when `proc_input` is `true`
577pub fn try_pfunc<'a, R>(
578    input: TokenStream,
579    proc_input: bool,
580    names: impl AsRef<[&'a str]>,
581    mut f: impl FnMut(Ident, Group) -> Result<TokenStream, R>,
582) -> Result<TokenStream, R> {
583    pfunc_impl(input, proc_input, names.as_ref(), &mut f)
584}
proc_macro_tool/lib.rs

proc_macro_tool/
lib.rs
