1use std::iter;
2
3use proc_macro2::TokenStream;
4use quote::TokenStreamExt;
5use quote::{quote, quote_spanned, ToTokens};
6use syn::visit_mut::VisitMut;
7use syn::{
8 punctuated::Punctuated, spanned::Spanned, Expr, ExprAsync, ExprCall, FieldPat, FnArg, Ident,
9 Item, ItemFn, Pat, PatIdent, PatReference, PatStruct, PatTuple, PatTupleStruct, PatType, Path,
10 ReturnType, Signature, Stmt, Token, Type, TypePath,
11};
12
13use crate::{
14 attr::{Field, FieldName, Fields, FormatMode, InstrumentArgs, Level},
15 MaybeItemFn, MaybeItemFnRef,
16};
17
18pub(crate) fn gen_function<'a, B: ToTokens + 'a>(
20 input: MaybeItemFnRef<'a, B>,
21 args: InstrumentArgs,
22 instrumented_function_name: &str,
23 self_type: Option<&TypePath>,
24) -> proc_macro2::TokenStream {
25 let MaybeItemFnRef {
29 outer_attrs,
30 inner_attrs,
31 vis,
32 sig,
33 brace_token,
34 block,
35 } = input;
36
37 let Signature {
38 output,
39 inputs: params,
40 unsafety,
41 asyncness,
42 constness,
43 abi,
44 ident,
45 generics:
46 syn::Generics {
47 params: gen_params,
48 where_clause,
49 lt_token,
50 gt_token,
51 },
52 fn_token,
53 paren_token,
54 variadic,
55 } = sig;
56
57 let warnings = args.warnings();
58
59 let (return_type, return_span) = if let ReturnType::Type(_, return_type) = &output {
60 (erase_impl_trait(return_type), return_type.span())
61 } else {
62 (syn::parse_quote! { () }, ident.span())
64 };
65 let fake_return_edge = quote_spanned! {return_span=>
72 #[allow(
73 unknown_lints,
74 unreachable_code,
75 clippy::diverging_sub_expression,
76 clippy::empty_loop,
77 clippy::let_unit_value,
78 clippy::let_with_type_underscore,
79 clippy::needless_return,
80 clippy::unreachable
81 )]
82 if false {
83 let __tracing_attr_fake_return: #return_type = loop {};
84 return __tracing_attr_fake_return;
85 }
86 };
87 let block = quote! {
88 {
89 #fake_return_edge
90 { #block }
91 }
92 };
93
94 let body = gen_block(
95 &block,
96 params,
97 asyncness.is_some(),
98 args,
99 instrumented_function_name,
100 self_type,
101 );
102
103 let mut result = quote!(
104 #(#outer_attrs) *
105 #vis #constness #asyncness #unsafety #abi #fn_token #ident
106 #lt_token #gen_params #gt_token
107 );
108
109 paren_token.surround(&mut result, |tokens| {
110 params.to_tokens(tokens);
111 variadic.to_tokens(tokens);
112 });
113
114 output.to_tokens(&mut result);
115 where_clause.to_tokens(&mut result);
116
117 brace_token.surround(&mut result, |tokens| {
118 tokens.append_all(inner_attrs);
119 warnings.to_tokens(tokens);
120 body.to_tokens(tokens);
121 });
122
123 result
124}
125
126fn gen_block<B: ToTokens>(
128 block: &B,
129 params: &Punctuated<FnArg, Token![,]>,
130 async_context: bool,
131 mut args: InstrumentArgs,
132 instrumented_function_name: &str,
133 self_type: Option<&TypePath>,
134) -> proc_macro2::TokenStream {
135 let span_name = args
137 .name
139 .as_ref()
140 .map(|name| quote!(#name))
141 .unwrap_or_else(|| quote!(#instrumented_function_name));
142
143 let args_level = args.level();
144 let level = args_level.clone();
145
146 let follows_from = args.follows_from.iter();
147 let follows_from = quote! {
148 #(for cause in #follows_from {
149 __tracing_attr_span.follows_from(cause);
150 })*
151 };
152
153 let span = (|| {
155 let param_names: Vec<(Ident, (Ident, RecordType))> = params
158 .clone()
159 .into_iter()
160 .flat_map(|param| match param {
161 FnArg::Typed(PatType { pat, ty, .. }) => {
162 param_names(*pat, RecordType::parse_from_ty(&ty))
163 }
164 FnArg::Receiver(_) => Box::new(iter::once((
165 Ident::new("self", param.span()),
166 RecordType::Debug,
167 ))),
168 })
169 .map(|(x, record_type)| {
180 if self_type.is_some() && x == "_self" {
183 (Ident::new("self", x.span()), (x, record_type))
184 } else {
185 (x.clone(), (x, record_type))
186 }
187 })
188 .collect();
189
190 for skip in &args.skips {
191 if !param_names.iter().map(|(user, _)| user).any(|y| y == skip) {
192 return quote_spanned! {skip.span()=>
193 compile_error!("attempting to skip non-existent parameter")
194 };
195 }
196 }
197
198 let target = args.target();
199
200 let parent = args.parent.iter();
201
202 let quoted_fields: Vec<_> = param_names
204 .iter()
205 .filter(|(param, _)| {
206 if args.skip_all || args.skips.contains(param) {
207 return false;
208 }
209
210 if let Some(ref fields) = args.fields {
213 fields.0.iter().all(|Field { ref name, .. }| {
214 match name {
215 FieldName::Expr(_) => true,
218 FieldName::Punctuated(punctuated) => {
219 let first = punctuated.first();
220 first != punctuated.last()
221 || !first.iter().any(|name| name == ¶m)
222 }
223 }
224 })
225 } else {
226 true
227 }
228 })
229 .map(|(user_name, (real_name, record_type))| match record_type {
230 RecordType::Value => quote!(#user_name = #real_name),
231 RecordType::Debug => quote!(#user_name = ::tracing::field::debug(&#real_name)),
232 })
233 .collect();
234
235 if let Some(Fields(ref mut fields)) = args.fields {
237 let mut replacer = IdentAndTypesRenamer {
238 idents: param_names.into_iter().map(|(a, (b, _))| (a, b)).collect(),
239 types: Vec::new(),
240 };
241
242 if let Some(self_type) = self_type {
245 replacer.types.push(("Self", self_type.clone()));
246 }
247
248 for e in fields.iter_mut().filter_map(|f| f.value.as_mut()) {
249 syn::visit_mut::visit_expr_mut(&mut replacer, e);
250 }
251 }
252
253 let custom_fields = &args.fields;
254
255 quote!(::tracing::span!(
256 target: #target,
257 #(parent: #parent,)*
258 #level,
259 #span_name,
260 #(#quoted_fields,)*
261 #custom_fields
262
263 ))
264 })();
265
266 let target = args.target();
267
268 let err_event = match args.err_args {
269 Some(event_args) => {
270 let level_tokens = event_args.level(Level::Error);
271 match event_args.mode {
272 FormatMode::Default | FormatMode::Display => Some(quote!(
273 ::tracing::event!(target: #target, #level_tokens, error = %e)
274 )),
275 FormatMode::Debug => Some(quote!(
276 ::tracing::event!(target: #target, #level_tokens, error = ?e)
277 )),
278 }
279 }
280 _ => None,
281 };
282
283 let ret_event = match args.ret_args {
284 Some(event_args) => {
285 let level_tokens = event_args.level(args_level);
286 match event_args.mode {
287 FormatMode::Display => Some(quote!(
288 ::tracing::event!(target: #target, #level_tokens, return = %x)
289 )),
290 FormatMode::Default | FormatMode::Debug => Some(quote!(
291 ::tracing::event!(target: #target, #level_tokens, return = ?x)
292 )),
293 }
294 }
295 _ => None,
296 };
297
298 if async_context {
306 let mk_fut = match (err_event, ret_event) {
307 (Some(err_event), Some(ret_event)) => quote_spanned!(block.span()=>
308 async move {
309 let __match_scrutinee = async move #block.await;
310 match __match_scrutinee {
311 #[allow(clippy::unit_arg)]
312 Ok(x) => {
313 #ret_event;
314 Ok(x)
315 },
316 Err(e) => {
317 #err_event;
318 Err(e)
319 }
320 }
321 }
322 ),
323 (Some(err_event), None) => quote_spanned!(block.span()=>
324 async move {
325 match async move #block.await {
326 #[allow(clippy::unit_arg)]
327 Ok(x) => Ok(x),
328 Err(e) => {
329 #err_event;
330 Err(e)
331 }
332 }
333 }
334 ),
335 (None, Some(ret_event)) => quote_spanned!(block.span()=>
336 async move {
337 let x = async move #block.await;
338 #ret_event;
339 x
340 }
341 ),
342 (None, None) => quote_spanned!(block.span()=>
343 async move #block
344 ),
345 };
346
347 return quote!(
348 let __tracing_attr_span = #span;
349 let __tracing_instrument_future = #mk_fut;
350 if !__tracing_attr_span.is_disabled() {
351 #follows_from
352 ::tracing::Instrument::instrument(
353 __tracing_instrument_future,
354 __tracing_attr_span
355 )
356 .await
357 } else {
358 __tracing_instrument_future.await
359 }
360 );
361 }
362
363 let span = quote!(
364 let __tracing_attr_span;
375 let __tracing_attr_guard;
376 if ::tracing::level_enabled!(#level) || ::tracing::if_log_enabled!(#level, {true} else {false}) {
377 __tracing_attr_span = #span;
378 #follows_from
379 __tracing_attr_guard = __tracing_attr_span.enter();
380 }
381 );
382
383 match (err_event, ret_event) {
384 (Some(err_event), Some(ret_event)) => quote_spanned! {block.span()=>
385 #span
386 #[allow(clippy::redundant_closure_call)]
387 match (move || #block)() {
388 #[allow(clippy::unit_arg)]
389 Ok(x) => {
390 #ret_event;
391 Ok(x)
392 },
393 Err(e) => {
394 #err_event;
395 Err(e)
396 }
397 }
398 },
399 (Some(err_event), None) => quote_spanned!(block.span()=>
400 #span
401 #[allow(clippy::redundant_closure_call)]
402 match (move || #block)() {
403 #[allow(clippy::unit_arg)]
404 Ok(x) => Ok(x),
405 Err(e) => {
406 #err_event;
407 Err(e)
408 }
409 }
410 ),
411 (None, Some(ret_event)) => quote_spanned!(block.span()=>
412 #span
413 #[allow(clippy::redundant_closure_call)]
414 let x = (move || #block)();
415 #ret_event;
416 x
417 ),
418 (None, None) => quote_spanned!(block.span() =>
419 #[allow(clippy::suspicious_else_formatting)]
424 {
425 #span
426 #[warn(clippy::suspicious_else_formatting)]
428 #block
429 }
430 ),
431 }
432}
433
434enum RecordType {
436 Value,
438 Debug,
440}
441
442impl RecordType {
443 const TYPES_FOR_VALUE: &'static [&'static str] = &[
445 "bool",
446 "str",
447 "u8",
448 "i8",
449 "u16",
450 "i16",
451 "u32",
452 "i32",
453 "u64",
454 "i64",
455 "u128",
456 "i128",
457 "f32",
458 "f64",
459 "usize",
460 "isize",
461 "String",
462 "NonZeroU8",
463 "NonZeroI8",
464 "NonZeroU16",
465 "NonZeroI16",
466 "NonZeroU32",
467 "NonZeroI32",
468 "NonZeroU64",
469 "NonZeroI64",
470 "NonZeroU128",
471 "NonZeroI128",
472 "NonZeroUsize",
473 "NonZeroIsize",
474 "Wrapping",
475 ];
476
477 fn parse_from_ty(ty: &Type) -> Self {
480 match ty {
481 Type::Path(TypePath { path, .. })
482 if path
483 .segments
484 .iter()
485 .next_back()
486 .map(|path_segment| {
487 let ident = path_segment.ident.to_string();
488 Self::TYPES_FOR_VALUE.iter().any(|&t| t == ident)
489 })
490 .unwrap_or(false) =>
491 {
492 RecordType::Value
493 }
494 Type::Reference(syn::TypeReference { elem, .. }) => RecordType::parse_from_ty(elem),
495 _ => RecordType::Debug,
496 }
497 }
498}
499
500fn param_names(pat: Pat, record_type: RecordType) -> Box<dyn Iterator<Item = (Ident, RecordType)>> {
501 match pat {
502 Pat::Ident(PatIdent { ident, .. }) => Box::new(iter::once((ident, record_type))),
503 Pat::Reference(PatReference { pat, .. }) => param_names(*pat, record_type),
504 Pat::Struct(PatStruct { fields, .. }) => Box::new(
509 fields
510 .into_iter()
511 .flat_map(|FieldPat { pat, .. }| param_names(*pat, RecordType::Debug)),
512 ),
513 Pat::Tuple(PatTuple { elems, .. }) => Box::new(
514 elems
515 .into_iter()
516 .flat_map(|p| param_names(p, RecordType::Debug)),
517 ),
518 Pat::TupleStruct(PatTupleStruct { elems, .. }) => Box::new(
519 elems
520 .into_iter()
521 .flat_map(|p| param_names(p, RecordType::Debug)),
522 ),
523
524 _ => Box::new(iter::empty()),
529 }
530}
531
532enum AsyncKind<'a> {
534 Function(&'a ItemFn),
537 Async {
541 async_expr: &'a ExprAsync,
542 pinned_box: bool,
543 },
544}
545
546pub(crate) struct AsyncInfo<'block> {
547 source_stmt: &'block Stmt,
549 kind: AsyncKind<'block>,
550 self_type: Option<TypePath>,
551 input: &'block ItemFn,
552}
553
554impl<'block> AsyncInfo<'block> {
555 pub(crate) fn from_fn(input: &'block ItemFn) -> Option<Self> {
581 if input.sig.asyncness.is_some() {
583 return None;
584 }
585
586 let block = &input.block;
587
588 let inside_funs = block.stmts.iter().filter_map(|stmt| {
590 if let Stmt::Item(Item::Fn(fun)) = &stmt {
591 if fun.sig.asyncness.is_some() {
593 return Some((stmt, fun));
594 }
595 }
596 None
597 });
598
599 let (last_expr_stmt, last_expr) = block.stmts.iter().rev().find_map(|stmt| {
602 if let Stmt::Expr(expr, _semi) = stmt {
603 Some((stmt, expr))
604 } else {
605 None
606 }
607 })?;
608
609 if let Expr::Async(async_expr) = last_expr {
611 return Some(AsyncInfo {
612 source_stmt: last_expr_stmt,
613 kind: AsyncKind::Async {
614 async_expr,
615 pinned_box: false,
616 },
617 self_type: None,
618 input,
619 });
620 }
621
622 let (outside_func, outside_args) = match last_expr {
624 Expr::Call(ExprCall { func, args, .. }) => (func, args),
625 _ => return None,
626 };
627
628 let path = match outside_func.as_ref() {
630 Expr::Path(path) => &path.path,
631 _ => return None,
632 };
633 if !path_to_string(path).ends_with("Box::pin") {
634 return None;
635 }
636
637 if outside_args.is_empty() {
641 return None;
642 }
643
644 if let Expr::Async(async_expr) = &outside_args[0] {
647 return Some(AsyncInfo {
648 source_stmt: last_expr_stmt,
649 kind: AsyncKind::Async {
650 async_expr,
651 pinned_box: true,
652 },
653 self_type: None,
654 input,
655 });
656 }
657
658 let func = match &outside_args[0] {
660 Expr::Call(ExprCall { func, .. }) => func,
661 _ => return None,
662 };
663
664 let func_name = match **func {
666 Expr::Path(ref func_path) => path_to_string(&func_path.path),
667 _ => return None,
668 };
669
670 let (stmt_func_declaration, func) = inside_funs
673 .into_iter()
674 .find(|(_, fun)| fun.sig.ident == func_name)?;
675
676 let mut self_type = None;
679 for arg in &func.sig.inputs {
680 if let FnArg::Typed(ty) = arg {
681 if let Pat::Ident(PatIdent { ref ident, .. }) = *ty.pat {
682 if ident == "_self" {
683 let mut ty = *ty.ty.clone();
684 if let Type::Reference(syn::TypeReference { elem, .. }) = ty {
686 ty = *elem;
687 }
688
689 if let Type::Path(tp) = ty {
690 self_type = Some(tp);
691 break;
692 }
693 }
694 }
695 }
696 }
697
698 Some(AsyncInfo {
699 source_stmt: stmt_func_declaration,
700 kind: AsyncKind::Function(func),
701 self_type,
702 input,
703 })
704 }
705
706 pub(crate) fn gen_async(
707 self,
708 args: InstrumentArgs,
709 instrumented_function_name: &str,
710 ) -> Result<proc_macro::TokenStream, syn::Error> {
711 let mut out_stmts: Vec<TokenStream> = self
713 .input
714 .block
715 .stmts
716 .iter()
717 .map(|stmt| stmt.to_token_stream())
718 .collect();
719
720 if let Some((iter, _stmt)) = self
721 .input
722 .block
723 .stmts
724 .iter()
725 .enumerate()
726 .find(|(_iter, stmt)| *stmt == self.source_stmt)
727 {
728 out_stmts[iter] = match self.kind {
730 AsyncKind::Function(fun) => {
732 let fun = MaybeItemFn::from(fun.clone());
733 gen_function(
734 fun.as_ref(),
735 args,
736 instrumented_function_name,
737 self.self_type.as_ref(),
738 )
739 }
740 AsyncKind::Async {
742 async_expr,
743 pinned_box,
744 } => {
745 let instrumented_block = gen_block(
746 &async_expr.block,
747 &self.input.sig.inputs,
748 true,
749 args,
750 instrumented_function_name,
751 None,
752 );
753 let async_attrs = &async_expr.attrs;
754 if pinned_box {
755 quote! {
756 ::std::boxed::Box::pin(#(#async_attrs) * async move { #instrumented_block })
757 }
758 } else {
759 quote! {
760 #(#async_attrs) * async move { #instrumented_block }
761 }
762 }
763 }
764 };
765 }
766
767 let vis = &self.input.vis;
768 let sig = &self.input.sig;
769 let attrs = &self.input.attrs;
770 Ok(quote!(
771 #(#attrs) *
772 #vis #sig {
773 #(#out_stmts) *
774 }
775 )
776 .into())
777 }
778}
779
780fn path_to_string(path: &Path) -> String {
782 use std::fmt::Write;
783 let mut res = String::with_capacity(path.segments.len() * 5);
785 for i in 0..path.segments.len() {
786 write!(&mut res, "{}", path.segments[i].ident)
787 .expect("writing to a String should never fail");
788 if i < path.segments.len() - 1 {
789 res.push_str("::");
790 }
791 }
792 res
793}
794
795struct IdentAndTypesRenamer<'a> {
800 types: Vec<(&'a str, TypePath)>,
801 idents: Vec<(Ident, Ident)>,
802}
803
804impl VisitMut for IdentAndTypesRenamer<'_> {
805 #[allow(clippy::cmp_owned)]
808 fn visit_ident_mut(&mut self, id: &mut Ident) {
809 for (old_ident, new_ident) in &self.idents {
810 if id.to_string() == old_ident.to_string() {
811 *id = new_ident.clone();
812 }
813 }
814 }
815
816 fn visit_type_mut(&mut self, ty: &mut Type) {
817 for (type_name, new_type) in &self.types {
818 if let Type::Path(TypePath { path, .. }) = ty {
819 if path_to_string(path) == *type_name {
820 *ty = Type::Path(new_type.clone());
821 }
822 }
823 }
824 }
825}
826
827struct ImplTraitEraser;
830
831impl VisitMut for ImplTraitEraser {
832 fn visit_type_mut(&mut self, t: &mut Type) {
833 if let Type::ImplTrait(..) = t {
834 *t = syn::TypeInfer {
835 underscore_token: Token),
836 }
837 .into();
838 } else {
839 syn::visit_mut::visit_type_mut(self, t);
840 }
841 }
842}
843
844fn erase_impl_trait(ty: &Type) -> Type {
845 let mut ty = ty.clone();
846 ImplTraitEraser.visit_type_mut(&mut ty);
847 ty
848}