From 3a36bb1b6f0c9b223f8d11741b0171ebe26edffe Mon Sep 17 00:00:00 2001
From: Rebecca Turner <rbt@sent.as>
Date: Wed, 16 Oct 2024 19:20:03 -0700
Subject: [PATCH] Let `dead_code` lint work on `#[instrument`ed functions

Closes (partially?) #1366
---
 tracing-attributes/src/expand.rs      | 37 +++++++++++++++++++--------
 tracing-attributes/src/lib.rs         | 15 +++++++++--
 tracing-attributes/tests/dead_code.rs | 10 ++++++++
 3 files changed, 49 insertions(+), 13 deletions(-)
 create mode 100644 tracing-attributes/tests/dead_code.rs

diff --git a/tracing-attributes/src/expand.rs b/tracing-attributes/src/expand.rs
index 37034e3ede..7276ba3a50 100644
--- a/tracing-attributes/src/expand.rs
+++ b/tracing-attributes/src/expand.rs
@@ -1,6 +1,7 @@
 use std::iter;
 
 use proc_macro2::TokenStream;
+use quote::TokenStreamExt;
 use quote::{quote, quote_spanned, ToTokens};
 use syn::visit_mut::VisitMut;
 use syn::{
@@ -29,6 +30,7 @@ pub(crate) fn gen_function<'a, B: ToTokens + 'a>(
         inner_attrs,
         vis,
         sig,
+        brace_token,
         block,
     } = input;
 
@@ -44,9 +46,12 @@ pub(crate) fn gen_function<'a, B: ToTokens + 'a>(
             syn::Generics {
                 params: gen_params,
                 where_clause,
-                ..
+                lt_token,
+                gt_token,
             },
-        ..
+        fn_token,
+        paren_token,
+        variadic,
     } = sig;
 
     let warnings = args.warnings();
@@ -90,16 +95,26 @@ pub(crate) fn gen_function<'a, B: ToTokens + 'a>(
         self_type,
     );
 
-    quote!(
+    let mut result = quote!(
         #(#outer_attrs) *
-        #vis #constness #asyncness #unsafety #abi fn #ident<#gen_params>(#params) #output
-        #where_clause
-        {
-            #(#inner_attrs) *
-            #warnings
-            #body
-        }
-    )
+        #vis #constness #asyncness #unsafety #abi #fn_token #ident
+        #lt_token #gen_params #gt_token
+    );
+
+    paren_token.surround(&mut result, |tokens| {
+        params.to_tokens(tokens);
+        variadic.to_tokens(tokens);
+    });
+
+    where_clause.to_tokens(&mut result);
+
+    brace_token.surround(&mut result, |tokens| {
+        tokens.append_all(inner_attrs);
+        warnings.to_tokens(tokens);
+        body.to_tokens(tokens);
+    });
+
+    result
 }
 
 /// Instrument a block
diff --git a/tracing-attributes/src/lib.rs b/tracing-attributes/src/lib.rs
index 809305ca4e..9a061be93c 100644
--- a/tracing-attributes/src/lib.rs
+++ b/tracing-attributes/src/lib.rs
@@ -80,8 +80,10 @@
 )]
 
 use proc_macro2::TokenStream;
+use quote::TokenStreamExt;
 use quote::{quote, ToTokens};
 use syn::parse::{Parse, ParseStream};
+use syn::token::Brace;
 use syn::{Attribute, ItemFn, Signature, Visibility};
 
 mod attr;
@@ -428,6 +430,7 @@ struct MaybeItemFn {
     inner_attrs: Vec<Attribute>,
     vis: Visibility,
     sig: Signature,
+    brace_token: Brace,
     block: TokenStream,
 }
 
@@ -438,6 +441,7 @@ impl MaybeItemFn {
             inner_attrs: &self.inner_attrs,
             vis: &self.vis,
             sig: &self.sig,
+            brace_token: &self.brace_token,
             block: &self.block,
         }
     }
@@ -451,12 +455,15 @@ impl Parse for MaybeItemFn {
         let vis: Visibility = input.parse()?;
         let sig: Signature = input.parse()?;
         let inner_attrs = input.call(Attribute::parse_inner)?;
-        let block: TokenStream = input.parse()?;
+        let block;
+        let brace_token = syn::braced!(block in input);
+        let block: TokenStream = block.call(|buffer| buffer.parse())?;
         Ok(Self {
             outer_attrs,
             inner_attrs,
             vis,
             sig,
+            brace_token,
             block,
         })
     }
@@ -474,12 +481,15 @@ impl From<ItemFn> for MaybeItemFn {
         let (outer_attrs, inner_attrs) = attrs
             .into_iter()
             .partition(|attr| attr.style == syn::AttrStyle::Outer);
+        let mut block_tokens = TokenStream::new();
+        block_tokens.append_all(block.stmts);
         Self {
             outer_attrs,
             inner_attrs,
             vis,
             sig,
-            block: block.to_token_stream(),
+            brace_token: block.brace_token,
+            block: block_tokens,
         }
     }
 }
@@ -492,5 +502,6 @@ struct MaybeItemFnRef<'a, B: ToTokens> {
     inner_attrs: &'a Vec<Attribute>,
     vis: &'a Visibility,
     sig: &'a Signature,
+    brace_token: &'a Brace,
     block: &'a B,
 }
diff --git a/tracing-attributes/tests/dead_code.rs b/tracing-attributes/tests/dead_code.rs
new file mode 100644
index 0000000000..6a3e05d4a2
--- /dev/null
+++ b/tracing-attributes/tests/dead_code.rs
@@ -0,0 +1,10 @@
+use tracing_attributes::instrument;
+
+#[deny(unfulfilled_lint_expectations)]
+#[expect(dead_code)]
+#[instrument]
+fn unused() {}
+
+#[expect(dead_code)]
+#[instrument]
+async fn unused_async() {}