extern crate proc_macro;
use proc_macro::TokenStream;
use quote::quote;
use syn::parenthesized;
use syn::parse::{Parse, ParseStream};
use syn::punctuated::Punctuated;
use syn::token::{Comma, Paren};
use syn::{parse_macro_input, Attribute, Expr, FnArg, Ident, Result, Token, Type, Visibility};
struct SingleLine {
attrs: Vec<Attribute>,
vis: Visibility,
_fn_token: Option<Token![fn]>,
async_token: Option<Token![async]>,
name: Ident,
args: Vec<FnArg>,
ret_ty: Option<Type>,
expr: Expr,
}
impl Parse for SingleLine {
fn parse(input: ParseStream) -> Result<Self> {
let attrs = input.call(Attribute::parse_outer)?;
let vis: Visibility = input.parse()?;
let async_token = if input.peek(Token!(async)) {
Some(input.parse()?)
} else {
None
};
let fn_token = if input.peek(Token![fn]) {
Some(input.parse()?)
} else {
None
};
let name: Ident = input.parse()?;
let mut args = Vec::new();
if input.peek(Paren) {
let content;
parenthesized!(content in input);
let punct: Punctuated<FnArg, Comma> = content.parse_terminated(FnArg::parse, Comma)?;
args.extend(punct);
}
let ret_ty = if input.peek(Token![->]) {
input.parse::<Token![->]>()?;
Some(input.parse()?)
} else {
None
};
input.parse::<Token![=>]>()?;
let expr: Expr = input.parse()?;
Ok(SingleLine {
attrs,
vis,
_fn_token: fn_token,
async_token,
name,
args,
ret_ty,
expr,
})
}
}
#[proc_macro]
pub fn single_line(item: TokenStream) -> TokenStream {
let input = parse_macro_input!(item as SingleLine);
let attrs = input.attrs;
let vis = input.vis;
let name = input.name;
let expr = input.expr;
let args = input.args;
let async_token = input.async_token;
let ret_ty = input.ret_ty.unwrap_or_else(|| syn::parse_quote! { () });
let is_method = args.iter().any(|arg| matches!(arg, FnArg::Receiver(_)));
let body = if is_method {
if let Expr::Path(ref path) = expr {
if path.path.segments.len() == 1 {
let field = &path.path.segments[0].ident;
quote! { self.#field }
} else {
quote! { #expr }
}
} else {
quote! { #expr }
}
} else {
quote! { #expr }
};
let args_quote = if !args.is_empty() {
quote! { ( #(#args),* ) }
} else {
quote! { () }
};
let expanded = quote! {
#(#attrs)*
#vis #async_token fn #name #args_quote -> #ret_ty {
#body
}
};
expanded.into()
}