| 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 | |