Support `#[track_caller]` on async closures

7 files changed, 77 insertions, 11 deletions

diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs
index 24e2985cf56..69192457ebd 100644
--- a/compiler/rustc_ast_lowering/src/expr.rs
+++ b/compiler/rustc_ast_lowering/src/expr.rs
@@ -31,6 +31,9 @@ impl<'hir> LoweringContext<'_, 'hir> {
 
     pub(super) fn lower_expr_mut(&mut self, e: &Expr) -> hir::Expr<'hir> {
         ensure_sufficient_stack(|| {
+            let hir_id = self.lower_node_id(e.id);
+            self.lower_attrs(hir_id, &e.attrs);
+
             let kind = match &e.kind {
                 ExprKind::Box(inner) => hir::ExprKind::Box(self.lower_expr(inner)),
                 ExprKind::Array(exprs) => hir::ExprKind::Array(self.lower_exprs(exprs)),
@@ -147,7 +150,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 ),
                 ExprKind::Async(capture_clause, closure_node_id, block) => self.make_async_expr(
                     *capture_clause,
-                    None,
+                    hir_id,
                     *closure_node_id,
                     None,
                     e.span,
@@ -184,6 +187,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                             binder,
                             *capture_clause,
                             e.id,
+                            hir_id,
                             *closure_id,
                             fn_decl,
                             body,
@@ -310,8 +314,6 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 ExprKind::MacCall(_) => panic!("{:?} shouldn't exist here", e.span),
             };
 
-            let hir_id = self.lower_node_id(e.id);
-            self.lower_attrs(hir_id, &e.attrs);
             hir::Expr { hir_id, kind, span: self.lower_span(e.span) }
         })
     }
@@ -576,7 +578,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
     pub(super) fn make_async_expr(
         &mut self,
         capture_clause: CaptureBy,
-        outer_hir_id: Option<hir::HirId>,
+        outer_hir_id: hir::HirId,
         closure_node_id: NodeId,
         ret_ty: Option<hir::FnRetTy<'hir>>,
         span: Span,
@@ -669,8 +671,9 @@ impl<'hir> LoweringContext<'_, 'hir> {
             hir::ExprKind::Closure(c)
         };
 
-        let track_caller = outer_hir_id
-            .and_then(|id| self.attrs.get(&id.local_id))
+        let track_caller = self
+            .attrs
+            .get(&outer_hir_id.local_id)
             .map_or(false, |attrs| attrs.into_iter().any(|attr| attr.has_name(sym::track_caller)));
 
         let hir_id = self.lower_node_id(closure_node_id);
@@ -985,6 +988,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         binder: &ClosureBinder,
         capture_clause: CaptureBy,
         closure_id: NodeId,
+        closure_hir_id: hir::HirId,
         inner_closure_id: NodeId,
         decl: &FnDecl,
         body: &Expr,
@@ -1018,9 +1022,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
 
                 let async_body = this.make_async_expr(
                     capture_clause,
-                    // FIXME(nbdd0121): This should also use a proper HIR id so `#[track_caller]`
-                    // can be applied on async closures as well.
-                    None,
+                    closure_hir_id,
                     inner_closure_id,
                     async_ret_ty,
                     body.span,
diff --git a/compiler/rustc_ast_lowering/src/item.rs b/compiler/rustc_ast_lowering/src/item.rs
index d73d6d3918e..73065ab5163 100644
--- a/compiler/rustc_ast_lowering/src/item.rs
+++ b/compiler/rustc_ast_lowering/src/item.rs
@@ -1139,7 +1139,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
 
             let async_expr = this.make_async_expr(
                 CaptureBy::Value,
-                Some(fn_id),
+                fn_id,
                 closure_id,
                 None,
                 body.span,
diff --git a/src/test/ui/async-await/track-caller/async-block.rs b/src/test/ui/async-await/track-caller/async-block.rs
new file mode 100644
index 00000000000..8e81387c34b
--- /dev/null
+++ b/src/test/ui/async-await/track-caller/async-block.rs
@@ -0,0 +1,9 @@
+// edition:2021
+
+#![feature(closure_track_caller, stmt_expr_attributes)]
+
+fn main() {
+    let _ = #[track_caller] async {
+        //~^ ERROR attribute should be applied to a function definition [E0739]
+    };
+}
diff --git a/src/test/ui/async-await/track-caller/async-block.stderr b/src/test/ui/async-await/track-caller/async-block.stderr
new file mode 100644
index 00000000000..407439921c0
--- /dev/null
+++ b/src/test/ui/async-await/track-caller/async-block.stderr
@@ -0,0 +1,12 @@
+error[E0739]: attribute should be applied to a function definition
+  --> $DIR/async-block.rs:6:13
+   |
+LL |       let _ = #[track_caller] async {
+   |  _____________^^^^^^^^^^^^^^^_-
+LL | |
+LL | |     };
+   | |_____- not a function definition
+
+error: aborting due to previous error
+
+For more information about this error, try `rustc --explain E0739`.
diff --git a/src/test/ui/async-await/track-caller/async-closure-gate.rs b/src/test/ui/async-await/track-caller/async-closure-gate.rs
new file mode 100644
index 00000000000..9593fdb1908
--- /dev/null
+++ b/src/test/ui/async-await/track-caller/async-closure-gate.rs
@@ -0,0 +1,10 @@
+// edition:2021
+
+#![feature(async_closure, stmt_expr_attributes)]
+
+fn main() {
+    let _ = #[track_caller] async || {
+        //~^ ERROR `#[track_caller]` on closures is currently unstable [E0658]
+        //~| ERROR `#[track_caller]` on closures is currently unstable [E0658]
+    };
+}
diff --git a/src/test/ui/async-await/track-caller/async-closure-gate.stderr b/src/test/ui/async-await/track-caller/async-closure-gate.stderr
new file mode 100644
index 00000000000..be3d110eccd
--- /dev/null
+++ b/src/test/ui/async-await/track-caller/async-closure-gate.stderr
@@ -0,0 +1,25 @@
+error[E0658]: `#[track_caller]` on closures is currently unstable
+  --> $DIR/async-closure-gate.rs:6:13
+   |
+LL |     let _ = #[track_caller] async || {
+   |             ^^^^^^^^^^^^^^^
+   |
+   = note: see issue #87417 <https://github.com/rust-lang/rust/issues/87417> for more information
+   = help: add `#![feature(closure_track_caller)]` to the crate attributes to enable
+
+error[E0658]: `#[track_caller]` on closures is currently unstable
+  --> $DIR/async-closure-gate.rs:6:38
+   |
+LL |       let _ = #[track_caller] async || {
+   |  ______________________________________^
+LL | |
+LL | |
+LL | |     };
+   | |_____^
+   |
+   = note: see issue #87417 <https://github.com/rust-lang/rust/issues/87417> for more information
+   = help: add `#![feature(closure_track_caller)]` to the crate attributes to enable
+
+error: aborting due to 2 previous errors
+
+For more information about this error, try `rustc --explain E0658`.
diff --git a/src/test/ui/async-await/track-caller/panic-track-caller.rs b/src/test/ui/async-await/track-caller/panic-track-caller.rs
index 5ebfeb3f36a..066cf97628f 100644
--- a/src/test/ui/async-await/track-caller/panic-track-caller.rs
+++ b/src/test/ui/async-await/track-caller/panic-track-caller.rs
@@ -1,7 +1,7 @@
 // run-pass
 // edition:2021
 // needs-unwind
-#![feature(closure_track_caller)]
+#![feature(closure_track_caller, async_closure, stmt_expr_attributes)]
 
 use std::future::Future;
 use std::panic;
@@ -67,6 +67,13 @@ async fn foo_assoc() {
     Foo::bar_assoc().await
 }
 
+async fn foo_closure() {
+    let c = #[track_caller] async || {
+        panic!();
+    };
+    c().await
+}
+
 fn panicked_at(f: impl FnOnce() + panic::UnwindSafe) -> u32 {
     let loc = Arc::new(Mutex::new(None));
 
@@ -87,4 +94,5 @@ fn main() {
     assert_eq!(panicked_at(|| block_on(foo())), 41);
     assert_eq!(panicked_at(|| block_on(foo_track_caller())), 54);
     assert_eq!(panicked_at(|| block_on(foo_assoc())), 67);
+    assert_eq!(panicked_at(|| block_on(foo_closure())), 74);
 }