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, Peekable},
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, $crate::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, $crate::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    + Sized
147{
148    fn push(&mut self, tt: TokenTree) -> &mut Self {
149        self.extend(once(tt));
150        self
151    }
152
153    fn add(&mut self, stream: TokenStream) -> &mut Self {
154        self.extend(once(stream));
155        self
156    }
157
158    fn parse_iter(self) -> ParseIter<Self::IntoIter> {
159        ParseIter { iter: self.into_iter(), buf: VecDeque::new() }
160    }
161
162    /// Split [`TokenStream`] to `predicate` false and true
163    fn split_with<F>(self, mut predicate: F) -> (Self, Peekable<Self::IntoIter>)
164    where F: FnMut(&TokenTree) -> bool,
165    {
166        let mut left = Self::default();
167
168        let mut iter = self.into_iter().peekable();
169        while let Some(tt) = iter
170            .next_if(|tt| ! predicate(tt))
171        {
172            left.push(tt);
173        }
174
175        (left, iter)
176    }
177}
178impl TokenStreamExt for TokenStream { }
179
180pub trait TokenTreeExt: Sized {
181    fn as_ident(&self) -> Option<&Ident>;
182    fn as_punct(&self) -> Option<&Punct>;
183    fn as_group(&self) -> Option<&Group>;
184    fn as_literal(&self) -> Option<&Literal>;
185    fn into_ident(self) -> Result<Ident, Self>;
186    fn into_punct(self) -> Result<Punct, Self>;
187    fn into_group(self) -> Result<Group, Self>;
188    fn into_literal(self) -> Result<Literal, Self>;
189
190    fn is_ident(&self) -> bool {
191        self.as_ident().is_some()
192    }
193
194    fn is_punct(&self) -> bool {
195        self.as_punct().is_some()
196    }
197
198    fn is_group(&self) -> bool {
199        self.as_group().is_some()
200    }
201
202    fn is_literal(&self) -> bool {
203        self.as_literal().is_some()
204    }
205
206    /// Ident content equal to `keyword` str
207    ///
208    /// Other return `false` when `self` is not [`Ident`]
209    fn is_keyword(&self, keyword: &str) -> bool {
210        self.as_ident().is_some_and(|i| i.to_string() == keyword)
211    }
212
213    /// Punct char equal to `ch`
214    ///
215    /// Other return `false` when `self` is not [`Punct`]
216    fn is_punch(&self, ch: char) -> bool {
217        self.as_punct().is_some_and(|p| p.as_char() == ch)
218    }
219
220    /// Group delimiter is not [`Delimiter::None`]
221    ///
222    /// Other return `false` when `self` is not [`Group`]
223    fn is_solid_group(&self) -> bool {
224        self.as_group().is_some_and(|g| g.is_solid_group())
225    }
226
227    /// Punct spacing is [`Joint`]
228    ///
229    /// Other return `false` when `self` is not [`Punct`]
230    fn is_joint(&self) -> bool {
231        self.as_punct().is_some_and(|p| p.spacing() == Joint)
232    }
233
234    fn as_punct_char(&self) -> Option<char> {
235        self.as_punct().map(|p| p.as_char())
236    }
237}
238impl TokenTreeExt for TokenTree {
239    fn as_ident(&self) -> Option<&Ident> {
240        match self {
241            TokenTree::Ident(i) => Some(i),
242            _ => None,
243        }
244    }
245
246    fn as_punct(&self) -> Option<&Punct> {
247        match self {
248            TokenTree::Punct(i) => Some(i),
249            _ => None,
250        }
251    }
252
253    fn as_group(&self) -> Option<&Group> {
254        match self {
255            TokenTree::Group(i) => Some(i),
256            _ => None,
257        }
258    }
259
260    fn as_literal(&self) -> Option<&Literal> {
261        match self {
262            TokenTree::Literal(i) => Some(i),
263            _ => None,
264        }
265    }
266
267    fn into_ident(self) -> Result<Ident, Self> {
268        match self {
269            TokenTree::Ident(i) => Ok(i),
270            _ => Err(self),
271        }
272    }
273
274    fn into_punct(self) -> Result<Punct, Self> {
275        match self {
276            TokenTree::Punct(i) => Ok(i),
277            _ => Err(self),
278        }
279    }
280
281    fn into_group(self) -> Result<Group, Self> {
282        match self {
283            TokenTree::Group(i) => Ok(i),
284            _ => Err(self),
285        }
286    }
287
288    fn into_literal(self) -> Result<Literal, Self> {
289        match self {
290            TokenTree::Literal(i) => Ok(i),
291            _ => Err(self),
292        }
293    }
294}
295
296pub trait GroupExt {
297    /// Group delimiter is not [`Delimiter::None`]
298    fn is_solid_group(&self) -> bool;
299}
300impl GroupExt for Group {
301    fn is_solid_group(&self) -> bool {
302        self.delimiter() != Delimiter::None
303    }
304}
305
306pub struct ParseIter<I: Iterator<Item = TokenTree>> {
307    iter: I,
308    buf: VecDeque<TokenTree>,
309}
310
311impl<I: Iterator<Item = TokenTree>> ParseIter<I> {
312    pub fn peek(&mut self) -> Option<&TokenTree> {
313        self.peek_i(0)
314    }
315
316    pub fn next_if<F>(&mut self, f: F) -> Option<TokenTree>
317    where F: FnOnce(&TokenTree) -> bool,
318    {
319        let peek = self.peek()?;
320
321        if f(peek) {
322            self.next()
323        } else {
324            None
325        }
326    }
327
328    pub fn peek_i(&mut self, i: usize) -> Option<&TokenTree> {
329        for _ in self.buf.len()..=i {
330            self.buf.push_back(self.iter.next()?);
331        }
332        Some(&self.buf[i])
333    }
334
335    pub fn peek_is<F>(&mut self, f: F) -> bool
336    where F: FnOnce(&TokenTree) -> bool,
337    {
338        self.peek().is_some_and(f)
339    }
340
341    pub fn peek_i_is<F>(&mut self, i: usize, f: F) -> bool
342    where F: FnOnce(&TokenTree) -> bool,
343    {
344        self.peek_i(i).is_some_and(f)
345    }
346
347    /// Peek jointed puncts
348    pub fn peek_puncts(&mut self, puncts: impl AsRef<[u8]>) -> Option<
349        impl Iterator<Item = &TokenTree>
350    > {
351        let mut prev = None;
352
353        for (i, ch) in puncts.as_ref().iter()
354            .copied().map(char::from).enumerate()
355        {
356            if let Some(prev) = prev {
357                if prev == Alone { return None }
358            }
359
360            let tt = self.peek_i(i)?;
361            let TokenTree::Punct(p) = tt else { return None };
362            if p.as_char() != ch { return None }
363
364            prev = Some(p.spacing());
365        }
366        Some(self.buf.iter())
367    }
368
369    /// Next jointed puncts
370    pub fn next_puncts(&mut self, puncts: impl AsRef<[u8]>) -> Option<
371        impl Iterator<Item = TokenTree> + '_
372    > {
373        let _ = self.peek_puncts(puncts.as_ref())?;
374        Some(self.buf.drain(..puncts.as_ref().len()))
375    }
376}
377pub trait ParseIterExt: Iterator<Item = TokenTree> + Sized {
378    fn parse_iter(self) -> ParseIter<Self> {
379        ParseIter { iter: self, buf: VecDeque::new() }
380    }
381}
382impl<I: Iterator<Item = TokenTree>> ParseIterExt for I { }
383
384impl<I: Iterator<Item = TokenTree>> Iterator for ParseIter<I> {
385    type Item = TokenTree;
386
387    fn next(&mut self) -> Option<Self::Item> {
388        self.buf.pop_front()
389            .or_else(|| self.iter.next())
390    }
391
392    fn count(self) -> usize
393    where Self: Sized,
394    {
395        self.buf.len() + self.iter.count()
396    }
397
398    fn fold<B, F>(self, init: B, f: F) -> B
399    where Self: Sized,
400          F: FnMut(B, Self::Item) -> B,
401    {
402        self.buf.into_iter().chain(self.iter).fold(init, f)
403    }
404
405    fn size_hint(&self) -> (usize, Option<usize>) {
406        let (lo, hi) = self.iter.size_hint();
407        let lo = lo.saturating_add(self.buf.len());
408        let hi = hi.and_then(|hi| hi.checked_add(self.buf.len()));
409        (lo, hi)
410    }
411}
412impl<I: DoubleEndedIterator<Item = TokenTree>> DoubleEndedIterator for ParseIter<I> {
413    fn next_back(&mut self) -> Option<Self::Item> {
414        self.iter.next_back()
415            .or_else(|| self.buf.pop_back())
416    }
417
418    fn rfold<B, F>(self, init: B, f: F) -> B
419    where Self: Sized,
420          F: FnMut(B, Self::Item) -> B,
421    {
422        self.buf.into_iter().chain(self.iter).rfold(init, f)
423    }
424}
425impl<I: ExactSizeIterator<Item = TokenTree>> ExactSizeIterator for ParseIter<I> { }
426impl<I: FusedIterator<Item = TokenTree>> FusedIterator for ParseIter<I> { }
427
428fn pfunc_impl<F, R>(
429    input: TokenStream,
430    proc_input: bool,
431    names: &[&str],
432    f: &mut F,
433) -> Result<TokenStream, R>
434where F: FnMut(Ident, Group) -> Result<TokenStream, R>,
435{
436    let mut iter = input.into_iter().parse_iter();
437    let mut result = TokenStream::new();
438
439    while let Some(tt) = iter.next() {
440        match tt {
441            TokenTree::Punct(p)
442                if p.as_char() == '#'
443                && iter.peek_is(|i| i.as_ident()
444                    .is_some_and(|i| names.contains(&&*i.to_string())))
445                && iter.peek_i_is(1, |t| t.is_solid_group())
446                =>
447            {
448                let ident = iter.next().unwrap().into_ident().unwrap();
449                let mut group = iter.next().unwrap().into_group().unwrap();
450                if proc_input {
451                    let sub = pfunc_impl(
452                        group.stream(), proc_input, names, f)?;
453                    group = Group::new(group.delimiter(), sub)
454                        .set_spaned(group.span());
455                }
456                result.add(f(ident, group)?);
457            },
458            TokenTree::Group(g) => {
459                let sub = pfunc_impl(g.stream(), proc_input, names, f)?;
460                let tt = Group::new(g.delimiter(), sub);
461                result.push(tt.set_spaned(g.span()).into());
462            },
463            _ => _ = result.push(tt),
464        }
465    }
466
467    Ok(result)
468}
469
470/// Call `f` on `#name(...)` `#name[...]` etc, exclude [`Delimiter::None`]
471///
472/// Apply pfunc for `(...)` when `proc_input` is `true`
473pub fn pfunc<'a>(
474    input: TokenStream,
475    proc_input: bool,
476    names: impl AsRef<[&'a str]>,
477    mut f: impl FnMut(Ident, Group) -> TokenStream,
478) -> TokenStream {
479    let f = &mut |i, g| {
480        Ok::<_, ()>(f(i, g))
481    };
482    pfunc_impl(input, proc_input, names.as_ref(), f).unwrap()
483}
484
485/// Call `f` on `#name(...)` `#name[...]` etc, exclude [`Delimiter::None`]
486///
487/// Apply pfunc for `(...)` when `proc_input` is `true`
488pub fn try_pfunc<'a, R>(
489    input: TokenStream,
490    proc_input: bool,
491    names: impl AsRef<[&'a str]>,
492    mut f: impl FnMut(Ident, Group) -> Result<TokenStream, R>,
493) -> Result<TokenStream, R> {
494    pfunc_impl(input, proc_input, names.as_ref(), &mut f)
495}
proc_macro_tool/lib.rs

proc_macro_tool/
lib.rs
