Skip to content

Commit

Permalink
macros: Improve IDE support
Browse files Browse the repository at this point in the history
  • Loading branch information
@nurmohammed840
nurmohammed840 committed Nov 12, 2024
1 parent bb7ca75 commit d5f5c74
Showing 1 changed file with 33 additions and 128 deletions.
161 changes: 33 additions & 128 deletions tokio-macros/src/entry.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
use proc_macro2::{Span, TokenStream, TokenTree};
use quote::{quote, quote_spanned, ToTokens};
use proc_macro2::{Span, TokenStream};
use quote::{quote, ToTokens};
use syn::parse::{Parse, ParseStream, Parser};
use syn::{braced, Attribute, Ident, Path, Signature, Visibility};
use syn::token::Brace;
use syn::{Attribute, Ident, Path, Signature, Visibility};

// syn::AttributeArgs does not implement syn::Parse
type AttributeArgs = syn::punctuated::Punctuated<syn::Meta, syn::Token![,]>;
Expand Down Expand Up @@ -392,41 +393,28 @@ fn build_config(
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.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
// `syn::Error::new_spanned` does.
let start = last_stmt.next().map_or_else(Span::call_site, |t| t.span());
let end = last_stmt.last().map_or(start, |t| t.span());
(start, end)
};

let crate_path = config
.crate_name
.map(ToTokens::into_token_stream)
.unwrap_or_else(|| Ident::new("tokio", last_stmt_start_span).into_token_stream());
.unwrap_or_else(|| Ident::new("tokio", Span::call_site()).into_token_stream());

let mut rt = match config.flavor {
RuntimeFlavor::CurrentThread => quote_spanned! {last_stmt_start_span=>
RuntimeFlavor::CurrentThread => quote! {
#crate_path::runtime::Builder::new_current_thread()
},
RuntimeFlavor::Threaded => quote_spanned! {last_stmt_start_span=>
RuntimeFlavor::Threaded => quote! {
#crate_path::runtime::Builder::new_multi_thread()
},
};
if let Some(v) = config.worker_threads {
rt = quote_spanned! {last_stmt_start_span=> #rt.worker_threads(#v) };
rt = quote! { #rt.worker_threads(#v) };
}
if let Some(v) = config.start_paused {
rt = quote_spanned! {last_stmt_start_span=> #rt.start_paused(#v) };
rt = quote! { #rt.start_paused(#v) };
}
if let Some(v) = config.unhandled_panic {
let unhandled_panic = v.into_tokens(&crate_path);
rt = quote_spanned! {last_stmt_start_span=> #rt.unhandled_panic(#unhandled_panic) };
rt = quote! { #rt.unhandled_panic(#unhandled_panic) };
}

let generated_attrs = if is_test {
Expand All @@ -439,7 +427,7 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt

let body_ident = quote! { body };
// This explicit `return` is intentional. See tokio-rs/tokio#4636
let last_block = quote_spanned! {last_stmt_end_span=>
let last_block = quote! {
#[allow(clippy::expect_used, clippy::diverging_sub_expression, clippy::needless_return)]
{
return #rt
Expand All @@ -450,7 +438,7 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
}
};

let body = input.body();
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
Expand All @@ -461,7 +449,7 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
//
// We don't do this for the main function as it should only be used once so
// there will be no benefit.
let body = if is_test {
input.body = if is_test {
let output_type = match &input.sig.output {
// For functions with no return value syn doesn't print anything,
// but that doesn't work as `Output` for our boxed `Future`, so
Expand All @@ -479,8 +467,7 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
let body = async #body;
}
};

input.into_tokens(generated_attrs, body, last_block)
input.into_tokens(generated_attrs, last_block)
}

fn token_stream_with_error(mut tokens: TokenStream, error: syn::Error) -> TokenStream {
Expand Down Expand Up @@ -549,7 +536,8 @@ pub(crate) fn test(args: TokenStream, item: TokenStream, rt_multi_thread: bool)
Ok(it) => it,
Err(e) => return token_stream_with_error(item, e),
};
let config = if let Some(attr) = input.attrs().find(|attr| is_test_attribute(attr)) {

let config = if let Some(attr) = input.attrs.iter().find(|attr| is_test_attribute(attr)) {
let msg = "second test attribute is supplied, consider removing or changing the order of your test attributes";
Err(syn::Error::new_spanned(attr, msg))
} else {
Expand All @@ -565,126 +553,43 @@ pub(crate) fn test(args: TokenStream, item: TokenStream, rt_multi_thread: bool)
}

struct ItemFn {
outer_attrs: Vec<Attribute>,
attrs: Vec<Attribute>,
vis: Visibility,
sig: Signature,
brace_token: syn::token::Brace,
inner_attrs: Vec<Attribute>,
stmts: Vec<proc_macro2::TokenStream>,
body: proc_macro2::TokenStream,
}

impl ItemFn {
/// Access all attributes of the function item.
fn attrs(&self) -> impl Iterator<Item = &Attribute> {
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,
}
impl Parse for ItemFn {
fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
Ok(Self {
attrs: input.call(Attribute::parse_outer)?,
vis: input.parse()?,
sig: input.parse()?,
body: input.parse()?,
})
}
}

impl ItemFn {
/// Convert our local function item into a token stream.
fn into_tokens(
self,
generated_attrs: proc_macro2::TokenStream,
body: proc_macro2::TokenStream,
last_block: proc_macro2::TokenStream,
) -> TokenStream {
let mut tokens = proc_macro2::TokenStream::new();
let mut tokens = generated_attrs;
// 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 behavior of `syn`.
for mut attr in self.inner_attrs {
attr.style = syn::AttrStyle::Outer;
for attr in self.attrs {
attr.to_tokens(&mut tokens);
}

// Add generated macros at the end, so macros processed later are aware of them.
generated_attrs.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);
Brace::default().surround(&mut tokens, |tokens| {
// Note: To add `TokenStream`, we need to use `Some(self.body)` instead of `self.body`.
tokens.extend(Some(self.body));
last_block.to_tokens(tokens);
});

tokens
}
}

impl Parse for ItemFn {
#[inline]
fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
// 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::<Option<syn::Token![;]>>()? {
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::<TokenTree>()?]);
}

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);
}
});
}
}

0 comments on commit d5f5c74

Please sign in to comment.