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

proc_macro_tool/
lib.rs
