proc_macro_tool/
func_utils.rs

1use proc_macro::{
2    Delimiter, Group, Ident, Literal, Punct, Spacing::*, Span, TokenStream,
3    TokenTree,
4};
5
6use crate::{
7    GetSpan, ParseIterExt as _, SetSpan, TokenStreamExt as _,
8    TokenTreeExt as _,
9};
10
11/// Create [`TokenStream`] from
12/// [`IntoIterator<Item = TokenTree>`](IntoIterator)
13#[must_use]
14pub fn stream<I>(iter: I) -> TokenStream
15where I: IntoIterator<Item = TokenTree>,
16{
17    TokenStream::from_iter(iter)
18}
19
20/// Create [`TokenStream`] from
21/// [`IntoIterator<Item = TokenStream>`](IntoIterator)
22#[must_use]
23pub fn streams<I>(iter: I) -> TokenStream
24where I: IntoIterator<Item = TokenStream>,
25{
26    TokenStream::from_iter(iter)
27}
28
29fn pfunc_impl<F, R>(
30    stream: impl IntoIterator<Item = TokenTree>,
31    proc_input: bool,
32    names: &[&str],
33    f: &mut F,
34) -> Result<TokenStream, R>
35where F: FnMut(Ident, Group) -> Result<TokenStream, R>,
36{
37    let mut iter = stream.into_iter().parse_iter();
38    let mut result = TokenStream::new();
39
40    while let Some(tt) = iter.next() {
41        match tt {
42            TokenTree::Punct(p)
43                if p.as_char() == '#'
44                && iter.peek_is(|i| i.as_ident()
45                    .is_some_and(|i| names.contains(&&*i.to_string())))
46                && iter.peek_i_is(1, |t| t.is_solid_group())
47                =>
48            {
49                let ident = iter.next().unwrap().into_ident().unwrap();
50                let mut group = iter.next().unwrap().into_group().unwrap();
51                if proc_input {
52                    let sub = pfunc_impl(
53                        group.stream(),
54                        proc_input,
55                        names,
56                        f,
57                    )?;
58                    group = sub
59                        .grouped(group.delimiter())
60                        .set_spaned(group.span());
61                }
62                result.add(f(ident, group)?);
63            },
64            TokenTree::Group(g) => {
65                let sub = pfunc_impl(g.stream(), proc_input, names, f)?;
66                result.push(sub
67                    .grouped(g.delimiter())
68                    .set_spaned(g.span())
69                    .into());
70            },
71            _ => _ = result.push(tt),
72        }
73    }
74
75    Ok(result)
76}
77
78/// Call `f` on `#name(...)` `#name[...]` etc, exclude [`Delimiter::None`]
79///
80/// Apply pfunc for `(...)` when `proc_input` is `true`
81#[allow(clippy::missing_panics_doc)]
82pub fn pfunc<'a>(
83    stream: impl IntoIterator<Item = TokenTree>,
84    proc_input: bool,
85    names: impl AsRef<[&'a str]>,
86    mut f: impl FnMut(Ident, Group) -> TokenStream,
87) -> TokenStream {
88    let f = &mut |i, g| {
89        Ok::<_, ()>(f(i, g))
90    };
91    pfunc_impl(stream, proc_input, names.as_ref(), f).unwrap()
92}
93
94/// Call `f` on `#name(...)` `#name[...]` etc, exclude [`Delimiter::None`]
95///
96/// Apply pfunc for `(...)` when `proc_input` is `true`
97pub fn try_pfunc<'a, R>(
98    stream: impl IntoIterator<Item = TokenTree>,
99    proc_input: bool,
100    names: impl AsRef<[&'a str]>,
101    mut f: impl FnMut(Ident, Group) -> Result<TokenStream, R>,
102) -> Result<TokenStream, R> {
103    pfunc_impl(stream, proc_input, names.as_ref(), &mut f)
104}
105
106/// Make `compile_error! {"..."}`
107#[must_use]
108pub fn err(msg: &str, span: impl GetSpan) -> TokenStream {
109    let s = span_setter(span.span());
110
111    stream([
112        Punct::new(':', Joint).into(),
113        Punct::new(':', Joint).into(),
114        Ident::new("core", span.span()).into(),
115        Punct::new(':', Joint).into(),
116        Punct::new(':', Joint).into(),
117        Ident::new("compile_error", span.span()).into(),
118        Punct::new('!', Joint).into(),
119        Group::new(Delimiter::Brace, stream([
120            Literal::string(msg).into(),
121        ].map(s))).into(),
122    ].map(s))
123}
124
125/// Like [`err()`], but use [`Result`]
126///
127/// # Errors
128/// - always return [`Err`]
129pub fn rerr<T>(msg: &str, span: impl GetSpan) -> Result<T, TokenStream> {
130    Err(err(msg, span))
131}
132
133/// Make puncts, `spacing` is last punct spacing
134///
135/// - `"+-"` like `[Joint('+'), Joint('-')]`
136/// - `"+- "` like `[Joint('+'), Alone('-')]`
137/// - `"+ -"` like `[Alone('+'), Joint('-')]`
138pub fn puncts(puncts: impl AsRef<[u8]>) -> TokenStream {
139    puncts_spanned(puncts, Span::call_site())
140}
141
142/// Make puncts, `spacing` is last punct spacing
143///
144/// Like [`puncts`], but `.set_span(span)`
145pub fn puncts_spanned(puncts: impl AsRef<[u8]>, span: Span) -> TokenStream {
146    let puncts = puncts.as_ref().trim_ascii_start();
147    let iter = &mut puncts.iter().copied().peekable();
148    let mut result = TokenStream::new();
149
150    while let Some(ch) = iter.next() {
151        debug_assert!(! ch.is_ascii_whitespace());
152        let mut s = None;
153        while iter.next_if(u8::is_ascii_whitespace).is_some() {
154            s = Some(Alone)
155        }
156        let spacing = s.or(iter.peek().map(|_| Joint))
157            .unwrap_or(Joint);
158        let p = Punct::new(ch.into(), spacing);
159        result.push(p.set_spaned(span).into());
160    }
161
162    result
163}
164
165/// Generate a function, set input `TokenTree` span
166pub fn span_setter<T>(span: Span) -> impl Fn(T) -> T + Copy
167where T: SetSpan,
168{
169    move |tt| {
170        tt.set_spaned(span)
171    }
172}
173
174/// Like [`Group::new(Delimiter::Parenthesis, iter)`](Group::new)
175pub fn paren<I>(iter: I) -> Group
176where I: IntoIterator,
177      TokenStream: FromIterator<I::Item>,
178{
179    iter.into_iter()
180        .collect::<TokenStream>()
181        .grouped_paren()
182}
183
184/// Like [`Group::new(Delimiter::Brace, iter)`](Group::new)
185pub fn brace<I>(iter: I) -> Group
186where I: IntoIterator,
187      TokenStream: FromIterator<I::Item>,
188{
189    iter.into_iter()
190        .collect::<TokenStream>()
191        .grouped_brace()
192}
193
194/// Like [`Group::new(Delimiter::Bracket, iter)`](Group::new)
195pub fn bracket<I>(iter: I) -> Group
196where I: IntoIterator,
197      TokenStream: FromIterator<I::Item>,
198{
199    iter.into_iter()
200        .collect::<TokenStream>()
201        .grouped_bracket()
202}
203
204/// Like [`Group::new(Delimiter::None, iter)`](Group::new)
205pub fn none<I>(iter: I) -> Group
206where I: IntoIterator,
207      TokenStream: FromIterator<I::Item>,
208{
209    iter.into_iter()
210        .collect::<TokenStream>()
211        .grouped_none()
212}