diff --git a/tokio-macros/src/entry.rs b/tokio-macros/src/entry.rs index 3b124cff9..edac53029 100644 --- a/tokio-macros/src/entry.rs +++ b/tokio-macros/src/entry.rs @@ -1,7 +1,7 @@ -use proc_macro::TokenStream; -use proc_macro2::Span; +use proc_macro2::{Span, TokenStream, TokenTree}; use quote::{quote, quote_spanned, ToTokens}; -use syn::{parse::Parser, Ident, Path}; +use syn::parse::{Parse, ParseStream, Parser}; +use syn::{braced, Attribute, Ident, Path, Signature, Visibility}; // syn::AttributeArgs does not implement syn::Parse type AttributeArgs = syn::punctuated::Punctuated; @@ -230,7 +230,7 @@ fn parse_bool(bool: syn::Lit, span: Span, field: &str) -> Result TokenStream { +fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenStream { input.sig.asyncness = None; // If type mismatch occurs, the current rustc points to the last statement. let (last_stmt_start_span, last_stmt_end_span) = { - let mut last_stmt = input - .block - .stmts - .last() - .map(ToTokens::into_token_stream) - .unwrap_or_default() - .into_iter(); + let mut last_stmt = input.stmts.last().cloned().unwrap_or_default().into_iter(); + // `Span` on stable Rust has a limitation that only points to the first // token, not the whole tokens. We can work around this limitation by // using the first/last span of the tokens like @@ -373,10 +368,8 @@ fn parse_knobs(mut input: syn::ItemFn, is_test: bool, config: FinalConfig) -> To quote! {} }; - let body = &input.block; - let brace_token = input.block.brace_token; let body_ident = quote! { body }; - let block_expr = quote_spanned! {last_stmt_end_span=> + let last_block = quote_spanned! {last_stmt_end_span=> #[allow(clippy::expect_used, clippy::diverging_sub_expression)] { return #rt @@ -387,6 +380,8 @@ fn parse_knobs(mut input: syn::ItemFn, is_test: bool, config: FinalConfig) -> To } }; + let body = input.body(); + // For test functions pin the body to the stack and use `Pin<&mut dyn // Future>` to reduce the amount of `Runtime::block_on` (and related // functions) copies we generate during compilation due to the generic @@ -415,25 +410,11 @@ fn parse_knobs(mut input: syn::ItemFn, is_test: bool, config: FinalConfig) -> To } }; - input.block = syn::parse2(quote! { - { - #body - #block_expr - } - }) - .expect("Parsing failure"); - input.block.brace_token = brace_token; - - let result = quote! { - #header - #input - }; - - result.into() + input.into_tokens(header, body, last_block) } fn token_stream_with_error(mut tokens: TokenStream, error: syn::Error) -> TokenStream { - tokens.extend(TokenStream::from(error.into_compile_error())); + tokens.extend(error.into_compile_error()); tokens } @@ -442,7 +423,7 @@ pub(crate) fn main(args: TokenStream, item: TokenStream, rt_multi_thread: bool) // If any of the steps for this macro fail, we still want to expand to an item that is as close // to the expected output as possible. This helps out IDEs such that completions and other // related features keep working. - let input: syn::ItemFn = match syn::parse(item.clone()) { + let input: ItemFn = match syn::parse2(item.clone()) { Ok(it) => it, Err(e) => return token_stream_with_error(item, e), }; @@ -452,8 +433,8 @@ pub(crate) fn main(args: TokenStream, item: TokenStream, rt_multi_thread: bool) Err(syn::Error::new_spanned(&input.sig.ident, msg)) } else { AttributeArgs::parse_terminated - .parse(args) - .and_then(|args| build_config(input.clone(), args, false, rt_multi_thread)) + .parse2(args) + .and_then(|args| build_config(&input, args, false, rt_multi_thread)) }; match config { @@ -466,21 +447,17 @@ pub(crate) fn test(args: TokenStream, item: TokenStream, rt_multi_thread: bool) // If any of the steps for this macro fail, we still want to expand to an item that is as close // to the expected output as possible. This helps out IDEs such that completions and other // related features keep working. - let input: syn::ItemFn = match syn::parse(item.clone()) { + let input: ItemFn = match syn::parse2(item.clone()) { Ok(it) => it, Err(e) => return token_stream_with_error(item, e), }; - let config = if let Some(attr) = input - .attrs - .iter() - .find(|attr| attr.meta.path().is_ident("test")) - { + let config = if let Some(attr) = input.attrs().find(|attr| attr.meta.path().is_ident("test")) { let msg = "second test attribute is supplied"; Err(syn::Error::new_spanned(attr, msg)) } else { AttributeArgs::parse_terminated - .parse(args) - .and_then(|args| build_config(input.clone(), args, true, rt_multi_thread)) + .parse2(args) + .and_then(|args| build_config(&input, args, true, rt_multi_thread)) }; match config { @@ -488,3 +465,127 @@ pub(crate) fn test(args: TokenStream, item: TokenStream, rt_multi_thread: bool) Err(e) => token_stream_with_error(parse_knobs(input, true, DEFAULT_ERROR_CONFIG), e), } } + +struct ItemFn { + outer_attrs: Vec, + vis: Visibility, + sig: Signature, + brace_token: syn::token::Brace, + inner_attrs: Vec, + stmts: Vec, +} + +impl ItemFn { + /// Access all attributes of the function item. + fn attrs(&self) -> impl Iterator { + self.outer_attrs.iter().chain(self.inner_attrs.iter()) + } + + /// Get the body of the function item in a manner so that it can be + /// conveniently used with the `quote!` macro. + fn body(&self) -> Body<'_> { + Body { + brace_token: self.brace_token, + stmts: &self.stmts, + } + } + + /// Convert our local function item into a token stream. + fn into_tokens( + self, + header: proc_macro2::TokenStream, + body: proc_macro2::TokenStream, + last_block: proc_macro2::TokenStream, + ) -> TokenStream { + let mut tokens = proc_macro2::TokenStream::new(); + header.to_tokens(&mut tokens); + + // Outer attributes are simply streamed as-is. + for attr in self.outer_attrs { + attr.to_tokens(&mut tokens); + } + + // Inner attributes require extra care, since they're not supported on + // blocks (which is what we're expanded into) we instead lift them + // outside of the function. This matches the behaviour of `syn`. + for mut attr in self.inner_attrs { + attr.style = syn::AttrStyle::Outer; + attr.to_tokens(&mut tokens); + } + + self.vis.to_tokens(&mut tokens); + self.sig.to_tokens(&mut tokens); + + self.brace_token.surround(&mut tokens, |tokens| { + body.to_tokens(tokens); + last_block.to_tokens(tokens); + }); + + tokens + } +} + +impl Parse for ItemFn { + #[inline] + fn parse(input: ParseStream<'_>) -> syn::Result { + // This parse implementation has been largely lifted from `syn`, with + // the exception of: + // * We don't have access to the plumbing necessary to parse inner + // attributes in-place. + // * We do our own statements parsing to avoid recursively parsing + // entire statements and only look for the parts we're interested in. + + let outer_attrs = input.call(Attribute::parse_outer)?; + let vis: Visibility = input.parse()?; + let sig: Signature = input.parse()?; + + let content; + let brace_token = braced!(content in input); + let inner_attrs = Attribute::parse_inner(&content)?; + + let mut buf = proc_macro2::TokenStream::new(); + let mut stmts = Vec::new(); + + while !content.is_empty() { + if let Some(semi) = content.parse::>()? { + semi.to_tokens(&mut buf); + stmts.push(buf); + buf = proc_macro2::TokenStream::new(); + continue; + } + + // Parse a single token tree and extend our current buffer with it. + // This avoids parsing the entire content of the sub-tree. + buf.extend([content.parse::()?]); + } + + if !buf.is_empty() { + stmts.push(buf); + } + + Ok(Self { + outer_attrs, + vis, + sig, + brace_token, + inner_attrs, + stmts, + }) + } +} + +struct Body<'a> { + brace_token: syn::token::Brace, + // Statements, with terminating `;`. + stmts: &'a [TokenStream], +} + +impl ToTokens for Body<'_> { + fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) { + self.brace_token.surround(tokens, |tokens| { + for stmt in self.stmts { + stmt.to_tokens(tokens); + } + }) + } +} diff --git a/tokio-macros/src/lib.rs b/tokio-macros/src/lib.rs index 11bbbaec8..f56af5282 100644 --- a/tokio-macros/src/lib.rs +++ b/tokio-macros/src/lib.rs @@ -204,7 +204,7 @@ use proc_macro::TokenStream; #[proc_macro_attribute] #[cfg(not(test))] // Work around for rust-lang/rust#62127 pub fn main(args: TokenStream, item: TokenStream) -> TokenStream { - entry::main(args, item, true) + entry::main(args.into(), item.into(), true).into() } /// Marks async function to be executed by selected runtime. This macro helps set up a `Runtime` @@ -269,7 +269,7 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream { #[proc_macro_attribute] #[cfg(not(test))] // Work around for rust-lang/rust#62127 pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream { - entry::main(args, item, false) + entry::main(args.into(), item.into(), false).into() } /// Marks async function to be executed by runtime, suitable to test environment. @@ -427,7 +427,7 @@ pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream { /// ``` #[proc_macro_attribute] pub fn test(args: TokenStream, item: TokenStream) -> TokenStream { - entry::test(args, item, true) + entry::test(args.into(), item.into(), true).into() } /// Marks async function to be executed by runtime, suitable to test environment @@ -442,7 +442,7 @@ pub fn test(args: TokenStream, item: TokenStream) -> TokenStream { /// ``` #[proc_macro_attribute] pub fn test_rt(args: TokenStream, item: TokenStream) -> TokenStream { - entry::test(args, item, false) + entry::test(args.into(), item.into(), false).into() } /// Always fails with the error message below.