diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs index ec9c393502056..d2db121a9c261 100644 --- a/compiler/rustc_ast_lowering/src/expr.rs +++ b/compiler/rustc_ast_lowering/src/expr.rs @@ -642,15 +642,40 @@ impl<'hir> LoweringContext<'_, 'hir> { hir::ExprKind::Closure(c) }; - let generator = hir::Expr { - hir_id: self.lower_node_id(closure_node_id), - kind: generator_kind, - span: self.lower_span(span), + let parent_has_track_caller = self + .attrs + .values() + .find(|attrs| attrs.into_iter().find(|attr| attr.has_name(sym::track_caller)).is_some()) + .is_some(); + let unstable_span = + self.mark_span_with_reason(DesugaringKind::Async, span, self.allow_gen_future.clone()); + + let hir_id = if parent_has_track_caller { + let generator_hir_id = self.lower_node_id(closure_node_id); + self.lower_attrs( + generator_hir_id, + &[Attribute { + kind: AttrKind::Normal(ptr::P(NormalAttr { + item: AttrItem { + path: Path::from_ident(Ident::new(sym::track_caller, span)), + args: MacArgs::Empty, + tokens: None, + }, + tokens: None, + })), + id: self.tcx.sess.parse_sess.attr_id_generator.mk_attr_id(), + style: AttrStyle::Outer, + span: unstable_span, + }], + ); + generator_hir_id + } else { + self.lower_node_id(closure_node_id) }; + let generator = hir::Expr { hir_id, kind: generator_kind, span: self.lower_span(span) }; + // `future::from_generator`: - let unstable_span = - self.mark_span_with_reason(DesugaringKind::Async, span, self.allow_gen_future.clone()); let gen_future = self.expr_lang_item_path( unstable_span, hir::LangItem::FromGenerator, diff --git a/compiler/rustc_ast_lowering/src/item.rs b/compiler/rustc_ast_lowering/src/item.rs index 76316a574acb7..e07f8374fe3d9 100644 --- a/compiler/rustc_ast_lowering/src/item.rs +++ b/compiler/rustc_ast_lowering/src/item.rs @@ -86,7 +86,7 @@ impl<'a, 'hir> ItemLowerer<'a, 'hir> { impl_trait_defs: Vec::new(), impl_trait_bounds: Vec::new(), allow_try_trait: Some([sym::try_trait_v2, sym::yeet_desugar_details][..].into()), - allow_gen_future: Some([sym::gen_future][..].into()), + allow_gen_future: Some([sym::gen_future, sym::closure_track_caller][..].into()), allow_into_future: Some([sym::into_future][..].into()), generics_def_id_map: Default::default(), }; diff --git a/library/core/src/future/mod.rs b/library/core/src/future/mod.rs index 6487aa0885955..107cf92c1c0f7 100644 --- a/library/core/src/future/mod.rs +++ b/library/core/src/future/mod.rs @@ -82,6 +82,7 @@ where impl> Future for GenFuture { type Output = T::Return; + #[track_caller] fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { // SAFETY: Safe because we're !Unpin + !Drop, and this is just a field projection. let gen = unsafe { Pin::map_unchecked_mut(self, |s| &mut s.0) }; 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 new file mode 100644 index 0000000000000..4e659da9ee069 --- /dev/null +++ b/src/test/ui/async-await/track-caller/panic-track-caller.rs @@ -0,0 +1,75 @@ +// run-pass +// edition:2021 +#![feature(closure_track_caller)] + +use std::future::Future; +use std::panic; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll, Wake}; +use std::thread::{self, Thread}; + +/// A waker that wakes up the current thread when called. +struct ThreadWaker(Thread); + +impl Wake for ThreadWaker { + fn wake(self: Arc) { + self.0.unpark(); + } +} + +/// Run a future to completion on the current thread. +fn block_on(fut: impl Future) -> T { + // Pin the future so it can be polled. + let mut fut = Box::pin(fut); + + // Create a new context to be passed to the future. + let t = thread::current(); + let waker = Arc::new(ThreadWaker(t)).into(); + let mut cx = Context::from_waker(&waker); + + // Run the future to completion. + loop { + match fut.as_mut().poll(&mut cx) { + Poll::Ready(res) => return res, + Poll::Pending => thread::park(), + } + } +} + +async fn bar() { + panic!() +} + +async fn foo() { + bar().await +} + +#[track_caller] +async fn bar_track_caller() { + panic!() +} + +async fn foo_track_caller() { + bar_track_caller().await +} + +fn panicked_at(f: impl FnOnce() + panic::UnwindSafe) -> u32 { + let loc = Arc::new(Mutex::new(None)); + + let hook = panic::take_hook(); + { + let loc = loc.clone(); + panic::set_hook(Box::new(move |info| { + *loc.lock().unwrap() = info.location().map(|loc| loc.line()) + })); + } + panic::catch_unwind(f).unwrap_err(); + panic::set_hook(hook); + let x = loc.lock().unwrap().unwrap(); + x +} + +fn main() { + assert_eq!(panicked_at(|| block_on(foo())), 40); + assert_eq!(panicked_at(|| block_on(foo_track_caller())), 53); +}