From fa99cb82690d99c0df2018d37aedaed29e40e131 Mon Sep 17 00:00:00 2001
From: Bryan Garza <1396101+bryangarza@users.noreply.github.com>
Date: Tue, 8 Nov 2022 23:45:55 +0000
Subject: Allow and add `track_caller` to generators

This patch allows the usage of the `track_caller` annotation on
generators, as well as sets them conditionally if the parent also has
`track_caller` set.

Also add this annotation on the `GenFuture`'s `poll()` function.
---
 compiler/rustc_ast_lowering/src/expr.rs | 60 ++++++++++++++++++++-------------
 1 file changed, 37 insertions(+), 23 deletions(-)

(limited to 'compiler/rustc_ast_lowering/src/expr.rs')

diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs
index 61d91874e61..6f68f679cc0 100644
--- a/compiler/rustc_ast_lowering/src/expr.rs
+++ b/compiler/rustc_ast_lowering/src/expr.rs
@@ -617,33 +617,47 @@ impl<'hir> LoweringContext<'_, 'hir> {
 
             hir::ExprKind::Closure(c)
         };
-        let generator_hir_id = self.lower_node_id(closure_node_id);
-        // FIXME: only add track caller if the parent is track_caller
-        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,
+        let mut parent_has_track_caller = false;
+        for attrs in self.attrs.values() {
+            for attr in attrs.into_iter() {
+                if attr.has_name(sym::track_caller) {
+                    parent_has_track_caller = true;
+                    break;
+                }
+            }
+            if parent_has_track_caller {
+                break;
+            }
+        }
+        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,
-                    },
-                    tokens: None,
-                })),
-                id: self.tcx.sess.parse_sess.attr_id_generator.mk_attr_id(),
-                style: AttrStyle::Outer,
-                span,
-            }],
-        );
-        let generator = hir::Expr {
-            hir_id: generator_hir_id,
-            kind: generator_kind,
-            span: self.lower_span(span),
+                    })),
+                    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,
-- 
cgit 1.4.1-3-g733a5