1 | use std::fmt::Display; |
2 | |
3 | use proc_macro2::{TokenStream, TokenTree}; |
4 | use quote::{ToTokens, TokenStreamExt}; |
5 | use syn::parse::{Parse, ParseStream}; |
6 | use syn::{braced, bracketed, token, AttrStyle, Attribute, Signature, Token, Visibility}; |
7 | |
8 | pub fn token_stream_with_error(mut tokens: TokenStream, error: syn::Error) -> TokenStream { |
9 | tokens.extend(iter:error.into_compile_error()); |
10 | tokens |
11 | } |
12 | |
13 | pub fn error<A: ToTokens, T: Display>(s: &mut TokenStream, obj: A, msg: T) { |
14 | s.extend(iter:syn::Error::new_spanned(tokens:obj.into_token_stream(), message:msg).into_compile_error()) |
15 | } |
16 | |
17 | /// Function signature and body. |
18 | /// |
19 | /// Same as `syn`'s `ItemFn` except we keep the body as a TokenStream instead of |
20 | /// parsing it. This makes the macro not error if there's a syntax error in the body, |
21 | /// which helps IDE autocomplete work better. |
22 | #[derive (Debug, Clone)] |
23 | pub struct ItemFn { |
24 | pub attrs: Vec<Attribute>, |
25 | pub vis: Visibility, |
26 | pub sig: Signature, |
27 | pub brace_token: token::Brace, |
28 | pub body: TokenStream, |
29 | } |
30 | |
31 | impl Parse for ItemFn { |
32 | fn parse(input: ParseStream) -> syn::Result<Self> { |
33 | let mut attrs = input.call(Attribute::parse_outer)?; |
34 | let vis: Visibility = input.parse()?; |
35 | let sig: Signature = input.parse()?; |
36 | |
37 | let content; |
38 | let brace_token = braced!(content in input); |
39 | while content.peek(Token![#]) && content.peek2(Token![!]) { |
40 | let content2; |
41 | attrs.push(Attribute { |
42 | pound_token: content.parse()?, |
43 | style: AttrStyle::Inner(content.parse()?), |
44 | bracket_token: bracketed!(content2 in content), |
45 | meta: content2.parse()?, |
46 | }); |
47 | } |
48 | |
49 | let mut body = Vec::new(); |
50 | while !content.is_empty() { |
51 | body.push(content.parse::<TokenTree>()?); |
52 | } |
53 | let body = body.into_iter().collect(); |
54 | |
55 | Ok(ItemFn { |
56 | attrs, |
57 | vis, |
58 | sig, |
59 | brace_token, |
60 | body, |
61 | }) |
62 | } |
63 | } |
64 | |
65 | impl ToTokens for ItemFn { |
66 | fn to_tokens(&self, tokens: &mut TokenStream) { |
67 | tokens.append_all(self.attrs.iter().filter(|a: &&Attribute| matches!(a.style, AttrStyle::Outer))); |
68 | self.vis.to_tokens(tokens); |
69 | self.sig.to_tokens(tokens); |
70 | self.brace_token.surround(tokens, |tokens: &mut TokenStream| { |
71 | tokens.append_all(self.body.clone()); |
72 | }); |
73 | } |
74 | } |
75 | |