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, Block, Expr, ExprAsync, ExprCall, FieldPat, FnArg,
9 Ident, Item, ItemFn, Pat, PatIdent, PatReference, PatStruct, PatTuple, PatTupleStruct, PatType,
10 Path, ReturnType, Signature, Stmt, Token, Type, TypePath,
11};
12
13use crate::{
14 attr::{Field, 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 let first = name.first();
215 first != name.last() || !first.iter().any(|name| name == ¶m)
216 })
217 } else {
218 true
219 }
220 })
221 .map(|(user_name, (real_name, record_type))| match record_type {
222 RecordType::Value => quote!(#user_name = #real_name),
223 RecordType::Debug => quote!(#user_name = ::tracing::field::debug(&#real_name)),
224 })
225 .collect();
226
227 if let Some(Fields(ref mut fields)) = args.fields {
229 let mut replacer = IdentAndTypesRenamer {
230 idents: param_names.into_iter().map(|(a, (b, _))| (a, b)).collect(),
231 types: Vec::new(),
232 };
233
234 if let Some(self_type) = self_type {
237 replacer.types.push(("Self", self_type.clone()));
238 }
239
240 for e in fields.iter_mut().filter_map(|f| f.value.as_mut()) {
241 syn::visit_mut::visit_expr_mut(&mut replacer, e);
242 }
243 }
244
245 let custom_fields = &args.fields;
246
247 quote!(::tracing::span!(
248 target: #target,
249 #(parent: #parent,)*
250 #level,
251 #span_name,
252 #(#quoted_fields,)*
253 #custom_fields
254
255 ))
256 })();
257
258 let target = args.target();
259
260 let err_event = match args.err_args {
261 Some(event_args) => {
262 let level_tokens = event_args.level(Level::Error);
263 match event_args.mode {
264 FormatMode::Default | FormatMode::Display => Some(quote!(
265 ::tracing::event!(target: #target, #level_tokens, error = %e)
266 )),
267 FormatMode::Debug => Some(quote!(
268 ::tracing::event!(target: #target, #level_tokens, error = ?e)
269 )),
270 }
271 }
272 _ => None,
273 };
274
275 let ret_event = match args.ret_args {
276 Some(event_args) => {
277 let level_tokens = event_args.level(args_level);
278 match event_args.mode {
279 FormatMode::Display => Some(quote!(
280 ::tracing::event!(target: #target, #level_tokens, return = %x)
281 )),
282 FormatMode::Default | FormatMode::Debug => Some(quote!(
283 ::tracing::event!(target: #target, #level_tokens, return = ?x)
284 )),
285 }
286 }
287 _ => None,
288 };
289
290 if async_context {
298 let mk_fut = match (err_event, ret_event) {
299 (Some(err_event), Some(ret_event)) => quote_spanned!(block.span()=>
300 async move {
301 let __match_scrutinee = async move #block.await;
302 match __match_scrutinee {
303 #[allow(clippy::unit_arg)]
304 Ok(x) => {
305 #ret_event;
306 Ok(x)
307 },
308 Err(e) => {
309 #err_event;
310 Err(e)
311 }
312 }
313 }
314 ),
315 (Some(err_event), None) => quote_spanned!(block.span()=>
316 async move {
317 match async move #block.await {
318 #[allow(clippy::unit_arg)]
319 Ok(x) => Ok(x),
320 Err(e) => {
321 #err_event;
322 Err(e)
323 }
324 }
325 }
326 ),
327 (None, Some(ret_event)) => quote_spanned!(block.span()=>
328 async move {
329 let x = async move #block.await;
330 #ret_event;
331 x
332 }
333 ),
334 (None, None) => quote_spanned!(block.span()=>
335 async move #block
336 ),
337 };
338
339 return quote!(
340 let __tracing_attr_span = #span;
341 let __tracing_instrument_future = #mk_fut;
342 if !__tracing_attr_span.is_disabled() {
343 #follows_from
344 ::tracing::Instrument::instrument(
345 __tracing_instrument_future,
346 __tracing_attr_span
347 )
348 .await
349 } else {
350 __tracing_instrument_future.await
351 }
352 );
353 }
354
355 let span = quote!(
356 let __tracing_attr_span;
367 let __tracing_attr_guard;
368 if ::tracing::level_enabled!(#level) || ::tracing::if_log_enabled!(#level, {true} else {false}) {
369 __tracing_attr_span = #span;
370 #follows_from
371 __tracing_attr_guard = __tracing_attr_span.enter();
372 }
373 );
374
375 match (err_event, ret_event) {
376 (Some(err_event), Some(ret_event)) => quote_spanned! {block.span()=>
377 #span
378 #[allow(clippy::redundant_closure_call)]
379 match (move || #block)() {
380 #[allow(clippy::unit_arg)]
381 Ok(x) => {
382 #ret_event;
383 Ok(x)
384 },
385 Err(e) => {
386 #err_event;
387 Err(e)
388 }
389 }
390 },
391 (Some(err_event), None) => quote_spanned!(block.span()=>
392 #span
393 #[allow(clippy::redundant_closure_call)]
394 match (move || #block)() {
395 #[allow(clippy::unit_arg)]
396 Ok(x) => Ok(x),
397 Err(e) => {
398 #err_event;
399 Err(e)
400 }
401 }
402 ),
403 (None, Some(ret_event)) => quote_spanned!(block.span()=>
404 #span
405 #[allow(clippy::redundant_closure_call)]
406 let x = (move || #block)();
407 #ret_event;
408 x
409 ),
410 (None, None) => quote_spanned!(block.span() =>
411 #[allow(clippy::suspicious_else_formatting)]
416 {
417 #span
418 #[warn(clippy::suspicious_else_formatting)]
420 #block
421 }
422 ),
423 }
424}
425
426enum RecordType {
428 Value,
430 Debug,
432}
433
434impl RecordType {
435 const TYPES_FOR_VALUE: &'static [&'static str] = &[
437 "bool",
438 "str",
439 "u8",
440 "i8",
441 "u16",
442 "i16",
443 "u32",
444 "i32",
445 "u64",
446 "i64",
447 "u128",
448 "i128",
449 "f32",
450 "f64",
451 "usize",
452 "isize",
453 "String",
454 "NonZeroU8",
455 "NonZeroI8",
456 "NonZeroU16",
457 "NonZeroI16",
458 "NonZeroU32",
459 "NonZeroI32",
460 "NonZeroU64",
461 "NonZeroI64",
462 "NonZeroU128",
463 "NonZeroI128",
464 "NonZeroUsize",
465 "NonZeroIsize",
466 "Wrapping",
467 ];
468
469 fn parse_from_ty(ty: &Type) -> Self {
472 match ty {
473 Type::Path(TypePath { path, .. })
474 if path
475 .segments
476 .iter()
477 .next_back()
478 .map(|path_segment| {
479 let ident = path_segment.ident.to_string();
480 Self::TYPES_FOR_VALUE.iter().any(|&t| t == ident)
481 })
482 .unwrap_or(false) =>
483 {
484 RecordType::Value
485 }
486 Type::Reference(syn::TypeReference { elem, .. }) => RecordType::parse_from_ty(elem),
487 _ => RecordType::Debug,
488 }
489 }
490}
491
492fn param_names(pat: Pat, record_type: RecordType) -> Box<dyn Iterator<Item = (Ident, RecordType)>> {
493 match pat {
494 Pat::Ident(PatIdent { ident, .. }) => Box::new(iter::once((ident, record_type))),
495 Pat::Reference(PatReference { pat, .. }) => param_names(*pat, record_type),
496 Pat::Struct(PatStruct { fields, .. }) => Box::new(
501 fields
502 .into_iter()
503 .flat_map(|FieldPat { pat, .. }| param_names(*pat, RecordType::Debug)),
504 ),
505 Pat::Tuple(PatTuple { elems, .. }) => Box::new(
506 elems
507 .into_iter()
508 .flat_map(|p| param_names(p, RecordType::Debug)),
509 ),
510 Pat::TupleStruct(PatTupleStruct { elems, .. }) => Box::new(
511 elems
512 .into_iter()
513 .flat_map(|p| param_names(p, RecordType::Debug)),
514 ),
515
516 _ => Box::new(iter::empty()),
521 }
522}
523
524enum AsyncKind<'a> {
526 Function(&'a ItemFn),
529 Async {
533 async_expr: &'a ExprAsync,
534 pinned_box: bool,
535 },
536}
537
538pub(crate) struct AsyncInfo<'block> {
539 source_stmt: &'block Stmt,
541 kind: AsyncKind<'block>,
542 self_type: Option<TypePath>,
543 input: &'block ItemFn,
544}
545
546impl<'block> AsyncInfo<'block> {
547 pub(crate) fn from_fn(input: &'block ItemFn) -> Option<Self> {
573 if input.sig.asyncness.is_some() {
575 return None;
576 }
577
578 let block = &input.block;
579
580 let inside_funs = block.stmts.iter().filter_map(|stmt| {
582 if let Stmt::Item(Item::Fn(fun)) = &stmt {
583 if fun.sig.asyncness.is_some() {
585 return Some((stmt, fun));
586 }
587 }
588 None
589 });
590
591 let (last_expr_stmt, last_expr) = block.stmts.iter().rev().find_map(|stmt| {
594 if let Stmt::Expr(expr, _semi) = stmt {
595 Some((stmt, expr))
596 } else {
597 None
598 }
599 })?;
600
601 if let Expr::Async(async_expr) = last_expr {
603 return Some(AsyncInfo {
604 source_stmt: last_expr_stmt,
605 kind: AsyncKind::Async {
606 async_expr,
607 pinned_box: false,
608 },
609 self_type: None,
610 input,
611 });
612 }
613
614 let (outside_func, outside_args) = match last_expr {
616 Expr::Call(ExprCall { func, args, .. }) => (func, args),
617 _ => return None,
618 };
619
620 let path = match outside_func.as_ref() {
622 Expr::Path(path) => &path.path,
623 _ => return None,
624 };
625 if !path_to_string(path).ends_with("Box::pin") {
626 return None;
627 }
628
629 if outside_args.is_empty() {
633 return None;
634 }
635
636 if let Expr::Async(async_expr) = &outside_args[0] {
639 return Some(AsyncInfo {
640 source_stmt: last_expr_stmt,
641 kind: AsyncKind::Async {
642 async_expr,
643 pinned_box: true,
644 },
645 self_type: None,
646 input,
647 });
648 }
649
650 let func = match &outside_args[0] {
652 Expr::Call(ExprCall { func, .. }) => func,
653 _ => return None,
654 };
655
656 let func_name = match **func {
658 Expr::Path(ref func_path) => path_to_string(&func_path.path),
659 _ => return None,
660 };
661
662 let (stmt_func_declaration, func) = inside_funs
665 .into_iter()
666 .find(|(_, fun)| fun.sig.ident == func_name)?;
667
668 let mut self_type = None;
671 for arg in &func.sig.inputs {
672 if let FnArg::Typed(ty) = arg {
673 if let Pat::Ident(PatIdent { ref ident, .. }) = *ty.pat {
674 if ident == "_self" {
675 let mut ty = *ty.ty.clone();
676 if let Type::Reference(syn::TypeReference { elem, .. }) = ty {
678 ty = *elem;
679 }
680
681 if let Type::Path(tp) = ty {
682 self_type = Some(tp);
683 break;
684 }
685 }
686 }
687 }
688 }
689
690 Some(AsyncInfo {
691 source_stmt: stmt_func_declaration,
692 kind: AsyncKind::Function(func),
693 self_type,
694 input,
695 })
696 }
697
698 pub(crate) fn gen_async(
699 self,
700 args: InstrumentArgs,
701 instrumented_function_name: &str,
702 ) -> Result<proc_macro::TokenStream, syn::Error> {
703 let mut out_stmts: Vec<TokenStream> = self
705 .input
706 .block
707 .stmts
708 .iter()
709 .map(|stmt| stmt.to_token_stream())
710 .collect();
711
712 if let Some((iter, _stmt)) = self
713 .input
714 .block
715 .stmts
716 .iter()
717 .enumerate()
718 .find(|(_iter, stmt)| *stmt == self.source_stmt)
719 {
720 out_stmts[iter] = match self.kind {
722 AsyncKind::Function(fun) => {
724 let fun = MaybeItemFn::from(fun.clone());
725 gen_function(
726 fun.as_ref(),
727 args,
728 instrumented_function_name,
729 self.self_type.as_ref(),
730 )
731 }
732 AsyncKind::Async {
734 async_expr,
735 pinned_box,
736 } => {
737 let instrumented_block = gen_block(
738 &async_expr.block,
739 &self.input.sig.inputs,
740 true,
741 args,
742 instrumented_function_name,
743 None,
744 );
745 let async_attrs = &async_expr.attrs;
746 if pinned_box {
747 quote! {
748 ::std::boxed::Box::pin(#(#async_attrs) * async move { #instrumented_block })
749 }
750 } else {
751 quote! {
752 #(#async_attrs) * async move { #instrumented_block }
753 }
754 }
755 }
756 };
757 }
758
759 let vis = &self.input.vis;
760 let sig = &self.input.sig;
761 let attrs = &self.input.attrs;
762 Ok(quote!(
763 #(#attrs) *
764 #vis #sig {
765 #(#out_stmts) *
766 }
767 )
768 .into())
769 }
770}
771
772fn path_to_string(path: &Path) -> String {
774 use std::fmt::Write;
775 let mut res = String::with_capacity(path.segments.len() * 5);
777 for i in 0..path.segments.len() {
778 write!(&mut res, "{}", path.segments[i].ident)
779 .expect("writing to a String should never fail");
780 if i < path.segments.len() - 1 {
781 res.push_str("::");
782 }
783 }
784 res
785}
786
787struct IdentAndTypesRenamer<'a> {
792 types: Vec<(&'a str, TypePath)>,
793 idents: Vec<(Ident, Ident)>,
794}
795
796impl VisitMut for IdentAndTypesRenamer<'_> {
797 #[allow(clippy::cmp_owned)]
800 fn visit_ident_mut(&mut self, id: &mut Ident) {
801 for (old_ident, new_ident) in &self.idents {
802 if id.to_string() == old_ident.to_string() {
803 *id = new_ident.clone();
804 }
805 }
806 }
807
808 fn visit_type_mut(&mut self, ty: &mut Type) {
809 for (type_name, new_type) in &self.types {
810 if let Type::Path(TypePath { path, .. }) = ty {
811 if path_to_string(path) == *type_name {
812 *ty = Type::Path(new_type.clone());
813 }
814 }
815 }
816 }
817}
818
819struct AsyncTraitBlockReplacer<'a> {
821 block: &'a Block,
822 patched_block: Block,
823}
824
825impl VisitMut for AsyncTraitBlockReplacer<'_> {
826 fn visit_block_mut(&mut self, i: &mut Block) {
827 if i == self.block {
828 *i = self.patched_block.clone();
829 }
830 }
831}
832
833struct ImplTraitEraser;
836
837impl VisitMut for ImplTraitEraser {
838 fn visit_type_mut(&mut self, t: &mut Type) {
839 if let Type::ImplTrait(..) = t {
840 *t = syn::TypeInfer {
841 underscore_token: Token),
842 }
843 .into();
844 } else {
845 syn::visit_mut::visit_type_mut(self, t);
846 }
847 }
848}
849
850fn erase_impl_trait(ty: &Type) -> Type {
851 let mut ty = ty.clone();
852 ImplTraitEraser.visit_type_mut(&mut ty);
853 ty
854}