Rollup merge of #61941 - cramertj:no-more-yield-errors, r=centril

Preserve generator and yield source for error messages

Previously, error messages after HIR lowering all referred
to generators and yield, regardless of whether the original
source was a generator or an async/await body. This change
tracks the kind of each generator and yield source in order
to provide appropriately tailored error messages.

Fixes #60615.

10 files changed, 178 insertions, 100 deletions

diff --git a/src/librustc/cfg/construct.rs b/src/librustc/cfg/construct.rs
index 213e57a3b37..a7750edbb6f 100644
--- a/src/librustc/cfg/construct.rs
+++ b/src/librustc/cfg/construct.rs
@@ -330,7 +330,7 @@ impl<'a, 'tcx> CFGBuilder<'a, 'tcx> {
             hir::ExprKind::DropTemps(ref e) |
             hir::ExprKind::Unary(_, ref e) |
             hir::ExprKind::Field(ref e, _) |
-            hir::ExprKind::Yield(ref e) |
+            hir::ExprKind::Yield(ref e, _) |
             hir::ExprKind::Repeat(ref e, _) => {
                 self.straightline(expr, pred, Some(&**e).into_iter())
             }
diff --git a/src/librustc/hir/intravisit.rs b/src/librustc/hir/intravisit.rs
index e29a373ec3b..8d7e51d1ea6 100644
--- a/src/librustc/hir/intravisit.rs
+++ b/src/librustc/hir/intravisit.rs
@@ -1089,7 +1089,7 @@ pub fn walk_expr<'v, V: Visitor<'v>>(visitor: &mut V, expression: &'v Expr) {
                 visitor.visit_expr(expr)
             }
         }
-        ExprKind::Yield(ref subexpression) => {
+        ExprKind::Yield(ref subexpression, _) => {
             visitor.visit_expr(subexpression);
         }
         ExprKind::Lit(_) | ExprKind::Err => {}
diff --git a/src/librustc/hir/lowering.rs b/src/librustc/hir/lowering.rs
index d2b8c4729a4..4baa61a377b 100644
--- a/src/librustc/hir/lowering.rs
+++ b/src/librustc/hir/lowering.rs
@@ -95,8 +95,7 @@ pub struct LoweringContext<'a> {
 
     modules: BTreeMap<NodeId, hir::ModuleItems>,
 
-    is_generator: bool,
-    is_async_body: bool,
+    generator_kind: Option<hir::GeneratorKind>,
 
     /// Used to get the current `fn`'s def span to point to when using `await`
     /// outside of an `async fn`.
@@ -264,8 +263,7 @@ pub fn lower_crate(
         current_hir_id_owner: vec![(CRATE_DEF_INDEX, 0)],
         item_local_id_counters: Default::default(),
         node_id_to_hir_id: IndexVec::new(),
-        is_generator: false,
-        is_async_body: false,
+        generator_kind: None,
         current_item: None,
         lifetimes_to_define: Vec::new(),
         is_collecting_in_band_lifetimes: false,
@@ -795,18 +793,49 @@ impl<'a> LoweringContext<'a> {
         })
     }
 
-    fn record_body(&mut self, arguments: HirVec<hir::Arg>, value: hir::Expr) -> hir::BodyId {
-        if self.is_generator && self.is_async_body {
-            span_err!(
-                self.sess,
-                value.span,
-                E0727,
-                "`async` generators are not yet supported",
-            );
-            self.sess.abort_if_errors();
+    fn generator_movability_for_fn(
+        &mut self,
+        decl: &ast::FnDecl,
+        fn_decl_span: Span,
+        generator_kind: Option<hir::GeneratorKind>,
+        movability: Movability,
+    ) -> Option<hir::GeneratorMovability> {
+        match generator_kind {
+            Some(hir::GeneratorKind::Gen) =>  {
+                if !decl.inputs.is_empty() {
+                    span_err!(
+                        self.sess,
+                        fn_decl_span,
+                        E0628,
+                        "generators cannot have explicit arguments"
+                    );
+                    self.sess.abort_if_errors();
+                }
+                Some(match movability {
+                    Movability::Movable => hir::GeneratorMovability::Movable,
+                    Movability::Static => hir::GeneratorMovability::Static,
+                })
+            },
+            Some(hir::GeneratorKind::Async) => {
+                bug!("non-`async` closure body turned `async` during lowering");
+            },
+            None => {
+                if movability == Movability::Static {
+                    span_err!(
+                        self.sess,
+                        fn_decl_span,
+                        E0697,
+                        "closures cannot be static"
+                    );
+                }
+                None
+            },
         }
+    }
+
+    fn record_body(&mut self, arguments: HirVec<hir::Arg>, value: hir::Expr) -> hir::BodyId {
         let body = hir::Body {
-            is_generator: self.is_generator || self.is_async_body,
+            generator_kind: self.generator_kind,
             arguments,
             value,
         };
@@ -1143,7 +1172,7 @@ impl<'a> LoweringContext<'a> {
         };
         let decl = self.lower_fn_decl(&ast_decl, None, /* impl trait allowed */ false, None);
         let body_id = self.lower_fn_body(&ast_decl, |this| {
-            this.is_async_body = true;
+            this.generator_kind = Some(hir::GeneratorKind::Async);
             body(this)
         });
         let generator = hir::Expr {
@@ -1168,12 +1197,10 @@ impl<'a> LoweringContext<'a> {
         &mut self,
         f: impl FnOnce(&mut LoweringContext<'_>) -> (HirVec<hir::Arg>, hir::Expr),
     ) -> hir::BodyId {
-        let prev_is_generator = mem::replace(&mut self.is_generator, false);
-        let prev_is_async_body = mem::replace(&mut self.is_async_body, false);
+        let prev_gen_kind = self.generator_kind.take();
         let (arguments, result) = f(self);
         let body_id = self.record_body(arguments, result);
-        self.is_generator = prev_is_generator;
-        self.is_async_body = prev_is_async_body;
+        self.generator_kind = prev_gen_kind;
         body_id
     }
 
@@ -4476,37 +4503,18 @@ impl<'a> LoweringContext<'a> {
 
                     self.with_new_scopes(|this| {
                         this.current_item = Some(fn_decl_span);
-                        let mut is_generator = false;
+                        let mut generator_kind = None;
                         let body_id = this.lower_fn_body(decl, |this| {
                             let e = this.lower_expr(body);
-                            is_generator = this.is_generator;
+                            generator_kind = this.generator_kind;
                             e
                         });
-                        let generator_option = if is_generator {
-                            if !decl.inputs.is_empty() {
-                                span_err!(
-                                    this.sess,
-                                    fn_decl_span,
-                                    E0628,
-                                    "generators cannot have explicit arguments"
-                                );
-                                this.sess.abort_if_errors();
-                            }
-                            Some(match movability {
-                                Movability::Movable => hir::GeneratorMovability::Movable,
-                                Movability::Static => hir::GeneratorMovability::Static,
-                            })
-                        } else {
-                            if movability == Movability::Static {
-                                span_err!(
-                                    this.sess,
-                                    fn_decl_span,
-                                    E0697,
-                                    "closures cannot be static"
-                                );
-                            }
-                            None
-                        };
+                        let generator_option = this.generator_movability_for_fn(
+                            &decl,
+                            fn_decl_span,
+                            generator_kind,
+                            movability,
+                        );
                         hir::ExprKind::Closure(
                             this.lower_capture_clause(capture_clause),
                             fn_decl,
@@ -4678,12 +4686,26 @@ impl<'a> LoweringContext<'a> {
             }
 
             ExprKind::Yield(ref opt_expr) => {
-                self.is_generator = true;
+                match self.generator_kind {
+                    Some(hir::GeneratorKind::Gen) => {},
+                    Some(hir::GeneratorKind::Async) => {
+                        span_err!(
+                            self.sess,
+                            e.span,
+                            E0727,
+                            "`async` generators are not yet supported",
+                        );
+                        self.sess.abort_if_errors();
+                    },
+                    None => {
+                        self.generator_kind = Some(hir::GeneratorKind::Gen);
+                    }
+                }
                 let expr = opt_expr
                     .as_ref()
                     .map(|x| self.lower_expr(x))
                     .unwrap_or_else(|| self.expr_unit(e.span));
-                hir::ExprKind::Yield(P(expr))
+                hir::ExprKind::Yield(P(expr), hir::YieldSource::Yield)
             }
 
             ExprKind::Err => hir::ExprKind::Err,
@@ -5755,19 +5777,23 @@ impl<'a> LoweringContext<'a> {
         //         yield ();
         //     }
         // }
-        if !self.is_async_body {
-            let mut err = struct_span_err!(
-                self.sess,
-                await_span,
-                E0728,
-                "`await` is only allowed inside `async` functions and blocks"
-            );
-            err.span_label(await_span, "only allowed inside `async` functions and blocks");
-            if let Some(item_sp) = self.current_item {
-                err.span_label(item_sp, "this is not `async`");
+        match self.generator_kind {
+            Some(hir::GeneratorKind::Async) => {},
+            Some(hir::GeneratorKind::Gen) |
+            None => {
+                let mut err = struct_span_err!(
+                    self.sess,
+                    await_span,
+                    E0728,
+                    "`await` is only allowed inside `async` functions and blocks"
+                );
+                err.span_label(await_span, "only allowed inside `async` functions and blocks");
+                if let Some(item_sp) = self.current_item {
+                    err.span_label(item_sp, "this is not `async`");
+                }
+                err.emit();
+                return hir::ExprKind::Err;
             }
-            err.emit();
-            return hir::ExprKind::Err;
         }
         let span = self.mark_span_with_reason(
             CompilerDesugaringKind::Await,
@@ -5865,7 +5891,7 @@ impl<'a> LoweringContext<'a> {
             let unit = self.expr_unit(span);
             let yield_expr = P(self.expr(
                 span,
-                hir::ExprKind::Yield(P(unit)),
+                hir::ExprKind::Yield(P(unit), hir::YieldSource::Await),
                 ThinVec::new(),
             ));
             self.stmt(span, hir::StmtKind::Expr(yield_expr))
diff --git a/src/librustc/hir/mod.rs b/src/librustc/hir/mod.rs
index 0884a726a27..b063a6541e2 100644
--- a/src/librustc/hir/mod.rs
+++ b/src/librustc/hir/mod.rs
@@ -1306,7 +1306,7 @@ pub struct BodyId {
 ///
 /// - an `arguments` array containing the `(x, y)` pattern
 /// - a `value` containing the `x + y` expression (maybe wrapped in a block)
-/// - `is_generator` would be false
+/// - `generator_kind` would be `None`
 ///
 /// All bodies have an **owner**, which can be accessed via the HIR
 /// map using `body_owner_def_id()`.
@@ -1314,7 +1314,7 @@ pub struct BodyId {
 pub struct Body {
     pub arguments: HirVec<Arg>,
     pub value: Expr,
-    pub is_generator: bool,
+    pub generator_kind: Option<GeneratorKind>,
 }
 
 impl Body {
@@ -1325,6 +1325,26 @@ impl Body {
     }
 }
 
+/// The type of source expression that caused this generator to be created.
+// Not `IsAsync` because we want to eventually add support for `AsyncGen`
+#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, HashStable,
+         RustcEncodable, RustcDecodable, Hash, Debug, Copy)]
+pub enum GeneratorKind {
+    /// An `async` block or function.
+    Async,
+    /// A generator literal created via a `yield` inside a closure.
+    Gen,
+}
+
+impl fmt::Display for GeneratorKind {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(match self {
+            GeneratorKind::Async => "`async` object",
+            GeneratorKind::Gen => "generator",
+        })
+    }
+}
+
 #[derive(Copy, Clone, Debug)]
 pub enum BodyOwnerKind {
     /// Functions and methods.
@@ -1531,8 +1551,8 @@ pub enum ExprKind {
     ///
     /// The final span is the span of the argument block `|...|`.
     ///
-    /// This may also be a generator literal, indicated by the final boolean,
-    /// in that case there is an `GeneratorClause`.
+    /// This may also be a generator literal or an `async block` as indicated by the
+    /// `Option<GeneratorMovability>`.
     Closure(CaptureClause, P<FnDecl>, BodyId, Span, Option<GeneratorMovability>),
     /// A block (e.g., `'label: { ... }`).
     Block(P<Block>, Option<Label>),
@@ -1576,7 +1596,7 @@ pub enum ExprKind {
     Repeat(P<Expr>, AnonConst),
 
     /// A suspension point for generators (i.e., `yield <expr>`).
-    Yield(P<Expr>),
+    Yield(P<Expr>, YieldSource),
 
     /// A placeholder for an expression that wasn't syntactically well formed in some way.
     Err,
@@ -1668,12 +1688,12 @@ pub enum LoopIdError {
 
 impl fmt::Display for LoopIdError {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        fmt::Display::fmt(match *self {
+        f.write_str(match self {
             LoopIdError::OutsideLoopScope => "not inside loop scope",
             LoopIdError::UnlabeledCfInWhileCondition =>
                 "unlabeled control flow (break or continue) in while condition",
             LoopIdError::UnresolvedLabel => "label not found",
-        }, f)
+        })
     }
 }
 
@@ -1687,13 +1707,34 @@ pub struct Destination {
     pub target_id: Result<HirId, LoopIdError>,
 }
 
+/// Whether a generator contains self-references, causing it to be `!Unpin`.
 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, HashStable,
          RustcEncodable, RustcDecodable, Hash, Debug, Copy)]
 pub enum GeneratorMovability {
+    /// May contain self-references, `!Unpin`.
     Static,
+    /// Must not contain self-references, `Unpin`.
     Movable,
 }
 
+/// The yield kind that caused an `ExprKind::Yield`.
+#[derive(Debug, Copy, Clone, RustcEncodable, RustcDecodable, HashStable)]
+pub enum YieldSource {
+    /// An `<expr>.await`.
+    Await,
+    /// A plain `yield`.
+    Yield,
+}
+
+impl fmt::Display for YieldSource {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(match self {
+            YieldSource::Await => "`await`",
+            YieldSource::Yield => "`yield`",
+        })
+    }
+}
+
 #[derive(Clone, RustcEncodable, RustcDecodable, Debug, Copy, HashStable)]
 pub enum CaptureClause {
     CaptureByValue,
@@ -2058,11 +2099,10 @@ impl Defaultness {
 
 impl fmt::Display for Unsafety {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        fmt::Display::fmt(match *self {
-                              Unsafety::Normal => "normal",
-                              Unsafety::Unsafe => "unsafe",
-                          },
-                          f)
+        f.write_str(match self {
+            Unsafety::Normal => "normal",
+            Unsafety::Unsafe => "unsafe",
+        })
     }
 }
 
@@ -2076,10 +2116,10 @@ pub enum ImplPolarity {
 
 impl fmt::Debug for ImplPolarity {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        match *self {
-            ImplPolarity::Positive => "positive".fmt(f),
-            ImplPolarity::Negative => "negative".fmt(f),
-        }
+        f.write_str(match self {
+            ImplPolarity::Positive => "positive",
+            ImplPolarity::Negative => "negative",
+        })
     }
 }
 
diff --git a/src/librustc/hir/print.rs b/src/librustc/hir/print.rs
index 7b0a499fa5c..69ee84dfa43 100644
--- a/src/librustc/hir/print.rs
+++ b/src/librustc/hir/print.rs
@@ -1501,7 +1501,7 @@ impl<'a> State<'a> {
 
                 self.pclose()?;
             }
-            hir::ExprKind::Yield(ref expr) => {
+            hir::ExprKind::Yield(ref expr, _) => {
                 self.word_space("yield")?;
                 self.print_expr_maybe_paren(&expr, parser::PREC_JUMP)?;
             }
diff --git a/src/librustc/ich/impls_hir.rs b/src/librustc/ich/impls_hir.rs
index f2f909b6af1..30d76f240d1 100644
--- a/src/librustc/ich/impls_hir.rs
+++ b/src/librustc/ich/impls_hir.rs
@@ -335,15 +335,15 @@ impl<'a> HashStable<StableHashingContext<'a>> for hir::Body {
                                           hcx: &mut StableHashingContext<'a>,
                                           hasher: &mut StableHasher<W>) {
         let hir::Body {
-            ref arguments,
-            ref value,
-            is_generator,
-        } = *self;
+            arguments,
+            value,
+            generator_kind,
+        } = self;
 
         hcx.with_node_id_hashing_mode(NodeIdHashingMode::Ignore, |hcx| {
             arguments.hash_stable(hcx, hasher);
             value.hash_stable(hcx, hasher);
-            is_generator.hash_stable(hcx, hasher);
+            generator_kind.hash_stable(hcx, hasher);
         });
     }
 }
diff --git a/src/librustc/infer/error_reporting/need_type_info.rs b/src/librustc/infer/error_reporting/need_type_info.rs
index 362a680f53c..fe151bdec6a 100644
--- a/src/librustc/infer/error_reporting/need_type_info.rs
+++ b/src/librustc/infer/error_reporting/need_type_info.rs
@@ -227,16 +227,15 @@ impl<'a, 'tcx> InferCtxt<'a, 'tcx> {
 
     pub fn need_type_info_err_in_generator(
         &self,
+        kind: hir::GeneratorKind,
         span: Span,
         ty: Ty<'tcx>,
     ) -> DiagnosticBuilder<'tcx> {
         let ty = self.resolve_vars_if_possible(&ty);
         let name = self.extract_type_name(&ty, None);
-
-        let mut err = struct_span_err!(self.tcx.sess,
-                       span,
-                       E0698,
-                       "type inside generator must be known in this context");
+        let mut err = struct_span_err!(
+            self.tcx.sess, span, E0698, "type inside {} must be known in this context", kind,
+        );
         err.span_label(span, InferCtxt::missing_type_msg(&name));
         err
     }
diff --git a/src/librustc/middle/expr_use_visitor.rs b/src/librustc/middle/expr_use_visitor.rs
index 1c3eead90fa..086ddfd7e33 100644
--- a/src/librustc/middle/expr_use_visitor.rs
+++ b/src/librustc/middle/expr_use_visitor.rs
@@ -546,7 +546,7 @@ impl<'a, 'tcx> ExprUseVisitor<'a, 'tcx> {
                 self.consume_expr(&base);
             }
 
-            hir::ExprKind::Yield(ref value) => {
+            hir::ExprKind::Yield(ref value, _) => {
                 self.consume_expr(&value);
             }
         }
diff --git a/src/librustc/middle/liveness.rs b/src/librustc/middle/liveness.rs
index dc6fa066071..bf054d68b70 100644
--- a/src/librustc/middle/liveness.rs
+++ b/src/librustc/middle/liveness.rs
@@ -1218,7 +1218,7 @@ impl<'a, 'tcx> Liveness<'a, 'tcx> {
             hir::ExprKind::Type(ref e, _) |
             hir::ExprKind::DropTemps(ref e) |
             hir::ExprKind::Unary(_, ref e) |
-            hir::ExprKind::Yield(ref e) |
+            hir::ExprKind::Yield(ref e, _) |
             hir::ExprKind::Repeat(ref e, _) => {
                 self.propagate_through_expr(&e, succ)
             }
diff --git a/src/librustc/middle/region.rs b/src/librustc/middle/region.rs
index dfab8e36bf9..93cb6ab96f8 100644
--- a/src/librustc/middle/region.rs
+++ b/src/librustc/middle/region.rs
@@ -331,7 +331,7 @@ pub struct ScopeTree {
     /// The reason is that semantically, until the `box` expression returns,
     /// the values are still owned by their containing expressions. So
     /// we'll see that `&x`.
-    yield_in_scope: FxHashMap<Scope, (Span, usize)>,
+    yield_in_scope: FxHashMap<Scope, YieldData>,
 
     /// The number of visit_expr and visit_pat calls done in the body.
     /// Used to sanity check visit_expr/visit_pat call count when
@@ -339,6 +339,15 @@ pub struct ScopeTree {
     body_expr_count: FxHashMap<hir::BodyId, usize>,
 }
 
+#[derive(Debug, Copy, Clone, RustcEncodable, RustcDecodable, HashStable)]
+pub struct YieldData {
+    /// `Span` of the yield.
+    pub span: Span,
+    /// The number of expressions and patterns appearing before the `yield` in the body + 1.
+    pub expr_and_pat_count: usize,
+    pub source: hir::YieldSource,
+}
+
 #[derive(Debug, Copy, Clone)]
 pub struct Context {
     /// the root of the current region tree. This is typically the id
@@ -695,7 +704,7 @@ impl<'tcx> ScopeTree {
     /// returns `Some((span, expr_count))` with the span of a yield we found and
     /// the number of expressions and patterns appearing before the `yield` in the body + 1.
     /// If there a are multiple yields in a scope, the one with the highest number is returned.
-    pub fn yield_in_scope(&self, scope: Scope) -> Option<(Span, usize)> {
+    pub fn yield_in_scope(&self, scope: Scope) -> Option<YieldData> {
         self.yield_in_scope.get(&scope).cloned()
     }
 
@@ -706,14 +715,14 @@ impl<'tcx> ScopeTree {
                                    scope: Scope,
                                    expr_hir_id: hir::HirId,
                                    body: &'tcx hir::Body) -> Option<Span> {
-        self.yield_in_scope(scope).and_then(|(span, count)| {
+        self.yield_in_scope(scope).and_then(|YieldData { span, expr_and_pat_count, .. }| {
             let mut visitor = ExprLocatorVisitor {
                 hir_id: expr_hir_id,
                 result: None,
                 expr_and_pat_count: 0,
             };
             visitor.visit_body(body);
-            if count >= visitor.result.unwrap() {
+            if expr_and_pat_count >= visitor.result.unwrap() {
                 Some(span)
             } else {
                 None
@@ -953,12 +962,16 @@ fn resolve_expr<'tcx>(visitor: &mut RegionResolutionVisitor<'tcx>, expr: &'tcx h
 
     debug!("resolve_expr post-increment {}, expr = {:?}", visitor.expr_and_pat_count, expr);
 
-    if let hir::ExprKind::Yield(..) = expr.node {
+    if let hir::ExprKind::Yield(_, source) = &expr.node {
         // Mark this expr's scope and all parent scopes as containing `yield`.
         let mut scope = Scope { id: expr.hir_id.local_id, data: ScopeData::Node };
         loop {
-            visitor.scope_tree.yield_in_scope.insert(scope,
-                (expr.span, visitor.expr_and_pat_count));
+            let data = YieldData {
+                span: expr.span,
+                expr_and_pat_count: visitor.expr_and_pat_count,
+                source: *source,
+            };
+            visitor.scope_tree.yield_in_scope.insert(scope, data);
 
             // Keep traversing up while we can.
             match visitor.scope_tree.parent_map.get(&scope) {
@@ -1302,7 +1315,7 @@ impl<'tcx> Visitor<'tcx> for RegionResolutionVisitor<'tcx> {
             resolve_local(self, None, Some(&body.value));
         }
 
-        if body.is_generator {
+        if body.generator_kind.is_some() {
             self.scope_tree.body_expr_count.insert(body_id, self.expr_and_pat_count);
         }