tokio_macros/
entry.rs

1use proc_macro2::{Span, TokenStream, TokenTree};
2use quote::{quote, quote_spanned, ToTokens};
3use syn::parse::{Parse, ParseStream, Parser};
4use syn::{braced, Attribute, Ident, Path, Signature, Visibility};
5
6// syn::AttributeArgs does not implement syn::Parse
7type AttributeArgs = syn::punctuated::Punctuated<syn::Meta, syn::Token![,]>;
8
9#[derive(Clone, Copy, PartialEq)]
10enum RuntimeFlavor {
11    CurrentThread,
12    Threaded,
13    Local,
14}
15
16impl RuntimeFlavor {
17    fn from_str(s: &str) -> Result<RuntimeFlavor, String> {
18        match s {
19            "current_thread" => Ok(RuntimeFlavor::CurrentThread),
20            "multi_thread" => Ok(RuntimeFlavor::Threaded),
21            "local" => Ok(RuntimeFlavor::Local),
22            "single_thread" => Err("The single threaded runtime flavor is called `current_thread`.".to_string()),
23            "basic_scheduler" => Err("The `basic_scheduler` runtime flavor has been renamed to `current_thread`.".to_string()),
24            "threaded_scheduler" => Err("The `threaded_scheduler` runtime flavor has been renamed to `multi_thread`.".to_string()),
25            _ => Err(format!("No such runtime flavor `{s}`. The runtime flavors are `current_thread`, `local`, and `multi_thread`.")),
26        }
27    }
28}
29
30#[derive(Clone, Copy, PartialEq)]
31enum UnhandledPanic {
32    Ignore,
33    ShutdownRuntime,
34}
35
36impl UnhandledPanic {
37    fn from_str(s: &str) -> Result<UnhandledPanic, String> {
38        match s {
39            "ignore" => Ok(UnhandledPanic::Ignore),
40            "shutdown_runtime" => Ok(UnhandledPanic::ShutdownRuntime),
41            _ => Err(format!("No such unhandled panic behavior `{s}`. The unhandled panic behaviors are `ignore` and `shutdown_runtime`.")),
42        }
43    }
44
45    fn into_tokens(self, crate_path: &TokenStream) -> TokenStream {
46        match self {
47            UnhandledPanic::Ignore => quote! { #crate_path::runtime::UnhandledPanic::Ignore },
48            UnhandledPanic::ShutdownRuntime => {
49                quote! { #crate_path::runtime::UnhandledPanic::ShutdownRuntime }
50            }
51        }
52    }
53}
54
55struct FinalConfig {
56    flavor: RuntimeFlavor,
57    worker_threads: Option<usize>,
58    start_paused: Option<bool>,
59    crate_name: Option<Path>,
60    unhandled_panic: Option<UnhandledPanic>,
61}
62
63/// Config used in case of the attribute not being able to build a valid config
64const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig {
65    flavor: RuntimeFlavor::CurrentThread,
66    worker_threads: None,
67    start_paused: None,
68    crate_name: None,
69    unhandled_panic: None,
70};
71
72struct Configuration {
73    rt_multi_thread_available: bool,
74    default_flavor: RuntimeFlavor,
75    flavor: Option<RuntimeFlavor>,
76    worker_threads: Option<(usize, Span)>,
77    start_paused: Option<(bool, Span)>,
78    is_test: bool,
79    crate_name: Option<Path>,
80    unhandled_panic: Option<(UnhandledPanic, Span)>,
81}
82
83impl Configuration {
84    fn new(is_test: bool, rt_multi_thread: bool) -> Self {
85        Configuration {
86            rt_multi_thread_available: rt_multi_thread,
87            default_flavor: match is_test {
88                true => RuntimeFlavor::CurrentThread,
89                false => RuntimeFlavor::Threaded,
90            },
91            flavor: None,
92            worker_threads: None,
93            start_paused: None,
94            is_test,
95            crate_name: None,
96            unhandled_panic: None,
97        }
98    }
99
100    fn set_flavor(&mut self, runtime: syn::Lit, span: Span) -> Result<(), syn::Error> {
101        if self.flavor.is_some() {
102            return Err(syn::Error::new(span, "`flavor` set multiple times."));
103        }
104
105        let runtime_str = parse_string(runtime, span, "flavor")?;
106        let runtime =
107            RuntimeFlavor::from_str(&runtime_str).map_err(|err| syn::Error::new(span, err))?;
108        self.flavor = Some(runtime);
109        Ok(())
110    }
111
112    fn set_worker_threads(
113        &mut self,
114        worker_threads: syn::Lit,
115        span: Span,
116    ) -> Result<(), syn::Error> {
117        if self.worker_threads.is_some() {
118            return Err(syn::Error::new(
119                span,
120                "`worker_threads` set multiple times.",
121            ));
122        }
123
124        let worker_threads = parse_int(worker_threads, span, "worker_threads")?;
125        if worker_threads == 0 {
126            return Err(syn::Error::new(span, "`worker_threads` may not be 0."));
127        }
128        self.worker_threads = Some((worker_threads, span));
129        Ok(())
130    }
131
132    fn set_start_paused(&mut self, start_paused: syn::Lit, span: Span) -> Result<(), syn::Error> {
133        if self.start_paused.is_some() {
134            return Err(syn::Error::new(span, "`start_paused` set multiple times."));
135        }
136
137        let start_paused = parse_bool(start_paused, span, "start_paused")?;
138        self.start_paused = Some((start_paused, span));
139        Ok(())
140    }
141
142    fn set_crate_name(&mut self, name: syn::Lit, span: Span) -> Result<(), syn::Error> {
143        if self.crate_name.is_some() {
144            return Err(syn::Error::new(span, "`crate` set multiple times."));
145        }
146        let name_path = parse_path(name, span, "crate")?;
147        self.crate_name = Some(name_path);
148        Ok(())
149    }
150
151    fn set_unhandled_panic(
152        &mut self,
153        unhandled_panic: syn::Lit,
154        span: Span,
155    ) -> Result<(), syn::Error> {
156        if self.unhandled_panic.is_some() {
157            return Err(syn::Error::new(
158                span,
159                "`unhandled_panic` set multiple times.",
160            ));
161        }
162
163        let unhandled_panic = parse_string(unhandled_panic, span, "unhandled_panic")?;
164        let unhandled_panic =
165            UnhandledPanic::from_str(&unhandled_panic).map_err(|err| syn::Error::new(span, err))?;
166        self.unhandled_panic = Some((unhandled_panic, span));
167        Ok(())
168    }
169
170    fn macro_name(&self) -> &'static str {
171        if self.is_test {
172            "tokio::test"
173        } else {
174            "tokio::main"
175        }
176    }
177
178    fn build(&self) -> Result<FinalConfig, syn::Error> {
179        use RuntimeFlavor as F;
180
181        let flavor = self.flavor.unwrap_or(self.default_flavor);
182
183        let worker_threads = match (flavor, self.worker_threads) {
184            (F::CurrentThread | F::Local, Some((_, worker_threads_span))) => {
185                let msg = format!(
186                    "The `worker_threads` option requires the `multi_thread` runtime flavor. Use `#[{}(flavor = \"multi_thread\")]`",
187                    self.macro_name(),
188                );
189                return Err(syn::Error::new(worker_threads_span, msg));
190            }
191            (F::CurrentThread | F::Local, None) => None,
192            (F::Threaded, worker_threads) if self.rt_multi_thread_available => {
193                worker_threads.map(|(val, _span)| val)
194            }
195            (F::Threaded, _) => {
196                let msg = if self.flavor.is_none() {
197                    "The default runtime flavor is `multi_thread`, but the `rt-multi-thread` feature is disabled."
198                } else {
199                    "The runtime flavor `multi_thread` requires the `rt-multi-thread` feature."
200                };
201                return Err(syn::Error::new(Span::call_site(), msg));
202            }
203        };
204
205        let start_paused = match (flavor, self.start_paused) {
206            (F::Threaded, Some((_, start_paused_span))) => {
207                let msg = format!(
208                    "The `start_paused` option requires the `current_thread` runtime flavor. Use `#[{}(flavor = \"current_thread\")]`",
209                    self.macro_name(),
210                );
211                return Err(syn::Error::new(start_paused_span, msg));
212            }
213            (F::CurrentThread | F::Local, Some((start_paused, _))) => Some(start_paused),
214            (_, None) => None,
215        };
216
217        let unhandled_panic = match (flavor, self.unhandled_panic) {
218            (F::Threaded, Some((_, unhandled_panic_span))) => {
219                let msg = format!(
220                    "The `unhandled_panic` option requires the `current_thread` runtime flavor. Use `#[{}(flavor = \"current_thread\")]`",
221                    self.macro_name(),
222                );
223                return Err(syn::Error::new(unhandled_panic_span, msg));
224            }
225            (F::CurrentThread | F::Local, Some((unhandled_panic, _))) => Some(unhandled_panic),
226            (_, None) => None,
227        };
228
229        Ok(FinalConfig {
230            crate_name: self.crate_name.clone(),
231            flavor,
232            worker_threads,
233            start_paused,
234            unhandled_panic,
235        })
236    }
237}
238
239fn parse_int(int: syn::Lit, span: Span, field: &str) -> Result<usize, syn::Error> {
240    match int {
241        syn::Lit::Int(lit) => match lit.base10_parse::<usize>() {
242            Ok(value) => Ok(value),
243            Err(e) => Err(syn::Error::new(
244                span,
245                format!("Failed to parse value of `{field}` as integer: {e}"),
246            )),
247        },
248        _ => Err(syn::Error::new(
249            span,
250            format!("Failed to parse value of `{field}` as integer."),
251        )),
252    }
253}
254
255fn parse_string(int: syn::Lit, span: Span, field: &str) -> Result<String, syn::Error> {
256    match int {
257        syn::Lit::Str(s) => Ok(s.value()),
258        syn::Lit::Verbatim(s) => Ok(s.to_string()),
259        _ => Err(syn::Error::new(
260            span,
261            format!("Failed to parse value of `{field}` as string."),
262        )),
263    }
264}
265
266fn parse_path(lit: syn::Lit, span: Span, field: &str) -> Result<Path, syn::Error> {
267    match lit {
268        syn::Lit::Str(s) => {
269            let err = syn::Error::new(
270                span,
271                format!(
272                    "Failed to parse value of `{}` as path: \"{}\"",
273                    field,
274                    s.value()
275                ),
276            );
277            s.parse::<syn::Path>().map_err(|_| err.clone())
278        }
279        _ => Err(syn::Error::new(
280            span,
281            format!("Failed to parse value of `{field}` as path."),
282        )),
283    }
284}
285
286fn parse_bool(bool: syn::Lit, span: Span, field: &str) -> Result<bool, syn::Error> {
287    match bool {
288        syn::Lit::Bool(b) => Ok(b.value),
289        _ => Err(syn::Error::new(
290            span,
291            format!("Failed to parse value of `{field}` as bool."),
292        )),
293    }
294}
295
296fn contains_impl_trait(ty: &syn::Type) -> bool {
297    match ty {
298        syn::Type::ImplTrait(_) => true,
299        syn::Type::Array(t) => contains_impl_trait(&t.elem),
300        syn::Type::Ptr(t) => contains_impl_trait(&t.elem),
301        syn::Type::Reference(t) => contains_impl_trait(&t.elem),
302        syn::Type::Slice(t) => contains_impl_trait(&t.elem),
303        syn::Type::Tuple(t) => t.elems.iter().any(contains_impl_trait),
304        syn::Type::Paren(t) => contains_impl_trait(&t.elem),
305        syn::Type::Group(t) => contains_impl_trait(&t.elem),
306        syn::Type::Path(t) => match t.path.segments.last() {
307            Some(segment) => match &segment.arguments {
308                syn::PathArguments::AngleBracketed(args) => args.args.iter().any(|arg| match arg {
309                    syn::GenericArgument::Type(t) => contains_impl_trait(t),
310                    syn::GenericArgument::AssocType(t) => contains_impl_trait(&t.ty),
311                    _ => false,
312                }),
313                syn::PathArguments::Parenthesized(args) => {
314                    args.inputs.iter().any(contains_impl_trait)
315                        || matches!(&args.output, syn::ReturnType::Type(_, t) if contains_impl_trait(t))
316                }
317                syn::PathArguments::None => false,
318            },
319            None => false,
320        },
321        _ => false,
322    }
323}
324
325fn build_config(
326    input: &ItemFn,
327    args: AttributeArgs,
328    is_test: bool,
329    rt_multi_thread: bool,
330) -> Result<FinalConfig, syn::Error> {
331    if input.sig.asyncness.is_none() {
332        let msg = "the `async` keyword is missing from the function declaration";
333        return Err(syn::Error::new_spanned(input.sig.fn_token, msg));
334    }
335
336    let mut config = Configuration::new(is_test, rt_multi_thread);
337    let macro_name = config.macro_name();
338
339    for arg in args {
340        match arg {
341            syn::Meta::NameValue(namevalue) => {
342                let ident = namevalue
343                    .path
344                    .get_ident()
345                    .ok_or_else(|| {
346                        syn::Error::new_spanned(&namevalue, "Must have specified ident")
347                    })?
348                    .to_string()
349                    .to_lowercase();
350                let lit = match &namevalue.value {
351                    syn::Expr::Lit(syn::ExprLit { lit, .. }) => lit,
352                    expr => return Err(syn::Error::new_spanned(expr, "Must be a literal")),
353                };
354                match ident.as_str() {
355                    "worker_threads" => {
356                        config.set_worker_threads(lit.clone(), syn::spanned::Spanned::span(lit))?;
357                    }
358                    "flavor" => {
359                        config.set_flavor(lit.clone(), syn::spanned::Spanned::span(lit))?;
360                    }
361                    "start_paused" => {
362                        config.set_start_paused(lit.clone(), syn::spanned::Spanned::span(lit))?;
363                    }
364                    "core_threads" => {
365                        let msg = "Attribute `core_threads` is renamed to `worker_threads`";
366                        return Err(syn::Error::new_spanned(namevalue, msg));
367                    }
368                    "crate" => {
369                        config.set_crate_name(lit.clone(), syn::spanned::Spanned::span(lit))?;
370                    }
371                    "unhandled_panic" => {
372                        config
373                            .set_unhandled_panic(lit.clone(), syn::spanned::Spanned::span(lit))?;
374                    }
375                    name => {
376                        let msg = format!(
377                            "Unknown attribute {name} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`",
378                        );
379                        return Err(syn::Error::new_spanned(namevalue, msg));
380                    }
381                }
382            }
383            syn::Meta::Path(path) => {
384                let name = path
385                    .get_ident()
386                    .ok_or_else(|| syn::Error::new_spanned(&path, "Must have specified ident"))?
387                    .to_string()
388                    .to_lowercase();
389                let msg = match name.as_str() {
390                    "threaded_scheduler" | "multi_thread" => {
391                        format!(
392                            "Set the runtime flavor with #[{macro_name}(flavor = \"multi_thread\")]."
393                        )
394                    }
395                    "basic_scheduler" | "current_thread" | "single_threaded" => {
396                        format!(
397                            "Set the runtime flavor with #[{macro_name}(flavor = \"current_thread\")]."
398                        )
399                    }
400                    "flavor" | "worker_threads" | "start_paused" | "crate" | "unhandled_panic" => {
401                        format!("The `{name}` attribute requires an argument.")
402                    }
403                    name => {
404                        format!("Unknown attribute {name} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`.")
405                    }
406                };
407                return Err(syn::Error::new_spanned(path, msg));
408            }
409            other => {
410                return Err(syn::Error::new_spanned(
411                    other,
412                    "Unknown attribute inside the macro",
413                ));
414            }
415        }
416    }
417
418    config.build()
419}
420
421fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenStream {
422    input.sig.asyncness = None;
423
424    // If type mismatch occurs, the current rustc points to the last statement.
425    let (last_stmt_start_span, last_stmt_end_span) = {
426        let mut last_stmt = input.stmts.last().cloned().unwrap_or_default().into_iter();
427
428        // `Span` on stable Rust has a limitation that only points to the first
429        // token, not the whole tokens. We can work around this limitation by
430        // using the first/last span of the tokens like
431        // `syn::Error::new_spanned` does.
432        let start = last_stmt.next().map_or_else(Span::call_site, |t| t.span());
433        let end = last_stmt.last().map_or(start, |t| t.span());
434        (start, end)
435    };
436
437    let crate_path = config
438        .crate_name
439        .map(ToTokens::into_token_stream)
440        .unwrap_or_else(|| {
441            Ident::new("tokio", Span::call_site().located_at(last_stmt_start_span))
442                .into_token_stream()
443        });
444
445    let use_builder = quote_spanned! {Span::call_site().located_at(last_stmt_start_span)=>
446        use #crate_path::runtime::Builder;
447    };
448
449    let mut rt = match config.flavor {
450        RuntimeFlavor::CurrentThread | RuntimeFlavor::Local => {
451            quote_spanned! {last_stmt_start_span=>
452                Builder::new_current_thread()
453            }
454        }
455        RuntimeFlavor::Threaded => quote_spanned! {last_stmt_start_span=>
456            Builder::new_multi_thread()
457        },
458    };
459
460    let mut checks = vec![];
461    let mut errors = vec![];
462
463    let build = if let RuntimeFlavor::Local = config.flavor {
464        checks.push(quote! { tokio_unstable });
465        errors.push("The local runtime flavor is only available when `tokio_unstable` is set.");
466        quote_spanned! {last_stmt_start_span=> build_local(Default::default())}
467    } else {
468        quote_spanned! {last_stmt_start_span=> build()}
469    };
470
471    if let Some(v) = config.worker_threads {
472        rt = quote_spanned! {last_stmt_start_span=> #rt.worker_threads(#v) };
473    }
474    if let Some(v) = config.start_paused {
475        rt = quote_spanned! {last_stmt_start_span=> #rt.start_paused(#v) };
476    }
477    if let Some(v) = config.unhandled_panic {
478        let unhandled_panic = v.into_tokens(&crate_path);
479        rt = quote_spanned! {last_stmt_start_span=> #rt.unhandled_panic(#unhandled_panic) };
480    }
481
482    let generated_attrs = if is_test {
483        quote! {
484            #[::core::prelude::v1::test]
485        }
486    } else {
487        quote! {}
488    };
489
490    let do_checks: TokenStream = checks
491        .iter()
492        .zip(&errors)
493        .map(|(check, error)| {
494            quote! {
495                #[cfg(not(#check))]
496                compile_error!(#error);
497            }
498        })
499        .collect();
500
501    let body_ident = quote! { body };
502    // This explicit `return` is intentional. See tokio-rs/tokio#4636
503    let last_block = quote_spanned! {last_stmt_end_span=>
504        #do_checks
505
506        #[cfg(all(#(#checks),*))]
507        #[allow(clippy::expect_used, clippy::diverging_sub_expression, clippy::needless_return, clippy::unwrap_in_result)]
508        {
509            #use_builder
510
511            return #rt
512                .enable_all()
513                .#build
514                .expect("Failed building the Runtime")
515                .block_on(#body_ident);
516        }
517
518        #[cfg(not(all(#(#checks),*)))]
519        {
520            panic!("fell through checks")
521        }
522    };
523
524    let body = input.body();
525
526    // For test functions pin the body to the stack and use `Pin<&mut dyn
527    // Future>` to reduce the amount of `Runtime::block_on` (and related
528    // functions) copies we generate during compilation due to the generic
529    // parameter `F` (the future to block on). This could have an impact on
530    // performance, but because it's only for testing it's unlikely to be very
531    // large.
532    //
533    // We don't do this for the main function as it should only be used once so
534    // there will be no benefit.
535    let output_type = match &input.sig.output {
536        // For functions with no return value syn doesn't print anything,
537        // but that doesn't work as `Output` for our boxed `Future`, so
538        // default to `()` (the same type as the function output).
539        syn::ReturnType::Default => quote! { () },
540        syn::ReturnType::Type(_, ret_type) => quote! { #ret_type },
541    };
542
543    let body = if is_test {
544        quote! {
545            let body = async #body;
546            #crate_path::pin!(body);
547            let body: ::core::pin::Pin<&mut dyn ::core::future::Future<Output = #output_type>> = body;
548        }
549    } else {
550        // force typecheck without runtime overhead
551        let check_block = match &input.sig.output {
552            syn::ReturnType::Type(_, t)
553                if matches!(**t, syn::Type::Never(_)) || contains_impl_trait(t) =>
554            {
555                quote! {}
556            }
557            _ => quote! {
558                if false {
559                    let _: &dyn ::core::future::Future<Output = #output_type> = &body;
560                }
561            },
562        };
563
564        quote! {
565            let body = async #body;
566            // Compile-time assertion that the future's output matches the return type.
567            let body = {
568                #check_block
569                body
570            };
571        }
572    };
573
574    input.into_tokens(generated_attrs, body, last_block)
575}
576
577fn token_stream_with_error(mut tokens: TokenStream, error: syn::Error) -> TokenStream {
578    tokens.extend(error.into_compile_error());
579    tokens
580}
581
582pub(crate) fn main(args: TokenStream, item: TokenStream, rt_multi_thread: bool) -> TokenStream {
583    // If any of the steps for this macro fail, we still want to expand to an item that is as close
584    // to the expected output as possible. This helps out IDEs such that completions and other
585    // related features keep working.
586    let input: ItemFn = match syn::parse2(item.clone()) {
587        Ok(it) => it,
588        Err(e) => return token_stream_with_error(item, e),
589    };
590
591    let config = if input.sig.ident == "main" && !input.sig.inputs.is_empty() {
592        let msg = "the main function cannot accept arguments";
593        Err(syn::Error::new_spanned(&input.sig.ident, msg))
594    } else {
595        AttributeArgs::parse_terminated
596            .parse2(args)
597            .and_then(|args| build_config(&input, args, false, rt_multi_thread))
598    };
599
600    match config {
601        Ok(config) => parse_knobs(input, false, config),
602        Err(e) => token_stream_with_error(parse_knobs(input, false, DEFAULT_ERROR_CONFIG), e),
603    }
604}
605
606// Check whether given attribute is a test attribute of forms:
607// * `#[test]`
608// * `#[core::prelude::*::test]` or `#[::core::prelude::*::test]`
609// * `#[std::prelude::*::test]` or `#[::std::prelude::*::test]`
610fn is_test_attribute(attr: &Attribute) -> bool {
611    let path = match &attr.meta {
612        syn::Meta::Path(path) => path,
613        _ => return false,
614    };
615    let candidates = [
616        ["core", "prelude", "*", "test"],
617        ["std", "prelude", "*", "test"],
618    ];
619    if path.leading_colon.is_none()
620        && path.segments.len() == 1
621        && path.segments[0].arguments.is_none()
622        && path.segments[0].ident == "test"
623    {
624        return true;
625    } else if path.segments.len() != candidates[0].len() {
626        return false;
627    }
628    candidates.into_iter().any(|segments| {
629        path.segments.iter().zip(segments).all(|(segment, path)| {
630            segment.arguments.is_none() && (path == "*" || segment.ident == path)
631        })
632    })
633}
634
635pub(crate) fn test(args: TokenStream, item: TokenStream, rt_multi_thread: bool) -> TokenStream {
636    // If any of the steps for this macro fail, we still want to expand to an item that is as close
637    // to the expected output as possible. This helps out IDEs such that completions and other
638    // related features keep working.
639    let input: ItemFn = match syn::parse2(item.clone()) {
640        Ok(it) => it,
641        Err(e) => return token_stream_with_error(item, e),
642    };
643    let config = if let Some(attr) = input.attrs().find(|attr| is_test_attribute(attr)) {
644        let msg = "second test attribute is supplied, consider removing or changing the order of your test attributes";
645        Err(syn::Error::new_spanned(attr, msg))
646    } else {
647        AttributeArgs::parse_terminated
648            .parse2(args)
649            .and_then(|args| build_config(&input, args, true, rt_multi_thread))
650    };
651
652    match config {
653        Ok(config) => parse_knobs(input, true, config),
654        Err(e) => token_stream_with_error(parse_knobs(input, true, DEFAULT_ERROR_CONFIG), e),
655    }
656}
657
658struct ItemFn {
659    outer_attrs: Vec<Attribute>,
660    vis: Visibility,
661    sig: Signature,
662    brace_token: syn::token::Brace,
663    inner_attrs: Vec<Attribute>,
664    stmts: Vec<proc_macro2::TokenStream>,
665}
666
667impl ItemFn {
668    /// Access all attributes of the function item.
669    fn attrs(&self) -> impl Iterator<Item = &Attribute> {
670        self.outer_attrs.iter().chain(self.inner_attrs.iter())
671    }
672
673    /// Get the body of the function item in a manner so that it can be
674    /// conveniently used with the `quote!` macro.
675    fn body(&self) -> Body<'_> {
676        Body {
677            brace_token: self.brace_token,
678            stmts: &self.stmts,
679        }
680    }
681
682    /// Convert our local function item into a token stream.
683    fn into_tokens(
684        self,
685        generated_attrs: proc_macro2::TokenStream,
686        body: proc_macro2::TokenStream,
687        last_block: proc_macro2::TokenStream,
688    ) -> TokenStream {
689        let mut tokens = proc_macro2::TokenStream::new();
690        // Outer attributes are simply streamed as-is.
691        for attr in self.outer_attrs {
692            attr.to_tokens(&mut tokens);
693        }
694
695        // Inner attributes require extra care, since they're not supported on
696        // blocks (which is what we're expanded into) we instead lift them
697        // outside of the function. This matches the behavior of `syn`.
698        for mut attr in self.inner_attrs {
699            attr.style = syn::AttrStyle::Outer;
700            attr.to_tokens(&mut tokens);
701        }
702
703        // Add generated macros at the end, so macros processed later are aware of them.
704        generated_attrs.to_tokens(&mut tokens);
705
706        self.vis.to_tokens(&mut tokens);
707        self.sig.to_tokens(&mut tokens);
708
709        self.brace_token.surround(&mut tokens, |tokens| {
710            body.to_tokens(tokens);
711            last_block.to_tokens(tokens);
712        });
713
714        tokens
715    }
716}
717
718impl Parse for ItemFn {
719    #[inline]
720    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
721        // This parse implementation has been largely lifted from `syn`, with
722        // the exception of:
723        // * We don't have access to the plumbing necessary to parse inner
724        //   attributes in-place.
725        // * We do our own statements parsing to avoid recursively parsing
726        //   entire statements and only look for the parts we're interested in.
727
728        let outer_attrs = input.call(Attribute::parse_outer)?;
729        let vis: Visibility = input.parse()?;
730        let sig: Signature = input.parse()?;
731
732        let content;
733        let brace_token = braced!(content in input);
734        let inner_attrs = Attribute::parse_inner(&content)?;
735
736        let mut buf = proc_macro2::TokenStream::new();
737        let mut stmts = Vec::new();
738
739        while !content.is_empty() {
740            if let Some(semi) = content.parse::<Option<syn::Token![;]>>()? {
741                semi.to_tokens(&mut buf);
742                stmts.push(buf);
743                buf = proc_macro2::TokenStream::new();
744                continue;
745            }
746
747            // Parse a single token tree and extend our current buffer with it.
748            // This avoids parsing the entire content of the sub-tree.
749            buf.extend([content.parse::<TokenTree>()?]);
750        }
751
752        if !buf.is_empty() {
753            stmts.push(buf);
754        }
755
756        Ok(Self {
757            outer_attrs,
758            vis,
759            sig,
760            brace_token,
761            inner_attrs,
762            stmts,
763        })
764    }
765}
766
767struct Body<'a> {
768    brace_token: syn::token::Brace,
769    // Statements, with terminating `;`.
770    stmts: &'a [TokenStream],
771}
772
773impl ToTokens for Body<'_> {
774    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
775        self.brace_token.surround(tokens, |tokens| {
776            for stmt in self.stmts {
777                stmt.to_tokens(tokens);
778            }
779        });
780    }
781}
tokio_macros/entry.rs

tokio_macros/
entry.rs
