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

proc_macro_tool/
lib.rs
