Refactor YieldKind so postfix yield must have an expression

14 files changed, 65 insertions, 30 deletions

diff --git a/compiler/rustc_ast/src/ast.rs b/compiler/rustc_ast/src/ast.rs
index 9dcdd868343..5b7545b3396 100644
--- a/compiler/rustc_ast/src/ast.rs
+++ b/compiler/rustc_ast/src/ast.rs
@@ -1657,7 +1657,7 @@ pub enum ExprKind {
     Try(P<Expr>),
 
     /// A `yield`, with an optional value to be yielded.
-    Yield(Option<P<Expr>>, YieldKind),
+    Yield(YieldKind),
 
     /// A `do yeet` (aka `throw`/`fail`/`bail`/`raise`/whatever),
     /// with an optional value to be returned.
@@ -1904,12 +1904,41 @@ pub enum MatchKind {
 }
 
 /// The kind of yield expression
-#[derive(Clone, Copy, Encodable, Decodable, Debug, PartialEq)]
+#[derive(Clone, Encodable, Decodable, Debug)]
 pub enum YieldKind {
     /// yield expr { ... }
-    Prefix,
+    Prefix(Option<P<Expr>>),
     /// expr.yield { ... }
-    Postfix,
+    Postfix(P<Expr>),
+}
+
+impl YieldKind {
+    /// Returns the expression inside the yield expression, if any.
+    ///
+    /// For postfix yields, this is guaranteed to be `Some`.
+    pub const fn expr(&self) -> Option<&P<Expr>> {
+        match self {
+            YieldKind::Prefix(expr) => expr.as_ref(),
+            YieldKind::Postfix(expr) => Some(expr),
+        }
+    }
+
+    /// Returns a mutable reference to the expression being yielded, if any.
+    pub const fn expr_mut(&mut self) -> Option<&mut P<Expr>> {
+        match self {
+            YieldKind::Prefix(expr) => expr.as_mut(),
+            YieldKind::Postfix(expr) => Some(expr),
+        }
+    }
+
+    /// Returns true if both yields are prefix or both are postfix.
+    pub const fn same_kind(&self, other: &Self) -> bool {
+        match (self, other) {
+            (YieldKind::Prefix(_), YieldKind::Prefix(_)) => true,
+            (YieldKind::Postfix(_), YieldKind::Postfix(_)) => true,
+            _ => false,
+        }
+    }
 }
 
 /// A literal in a meta item.
diff --git a/compiler/rustc_ast/src/mut_visit.rs b/compiler/rustc_ast/src/mut_visit.rs
index 5a9df6ffadf..b159e136245 100644
--- a/compiler/rustc_ast/src/mut_visit.rs
+++ b/compiler/rustc_ast/src/mut_visit.rs
@@ -1813,8 +1813,11 @@ pub fn walk_expr<T: MutVisitor>(vis: &mut T, Expr { kind, id, span, attrs, token
         ExprKind::Paren(expr) => {
             vis.visit_expr(expr);
         }
-        ExprKind::Yield(expr, _) => {
-            visit_opt(expr, |expr| vis.visit_expr(expr));
+        ExprKind::Yield(kind) => {
+            let expr = kind.expr_mut();
+            if let Some(expr) = expr {
+                vis.visit_expr(expr);
+            }
         }
         ExprKind::Try(expr) => vis.visit_expr(expr),
         ExprKind::TryBlock(body) => vis.visit_block(body),
diff --git a/compiler/rustc_ast/src/util/classify.rs b/compiler/rustc_ast/src/util/classify.rs
index 116847af9f4..989ebe14bf8 100644
--- a/compiler/rustc_ast/src/util/classify.rs
+++ b/compiler/rustc_ast/src/util/classify.rs
@@ -182,11 +182,14 @@ pub fn expr_trailing_brace(mut expr: &ast::Expr) -> Option<TrailingBrace<'_>> {
             | Range(_, Some(e), _)
             | Ret(Some(e))
             | Unary(_, e)
-            | Yield(Some(e), _)
             | Yeet(Some(e))
             | Become(e) => {
                 expr = e;
             }
+            Yield(kind) => match kind.expr() {
+                Some(e) => expr = e,
+                None => break None,
+            },
             Closure(closure) => {
                 expr = &closure.body;
             }
@@ -217,7 +220,6 @@ pub fn expr_trailing_brace(mut expr: &ast::Expr) -> Option<TrailingBrace<'_>> {
             Break(_, None)
             | Range(_, None, _)
             | Ret(None)
-            | Yield(None, _)
             | Array(_)
             | Call(_, _)
             | MethodCall(_)
@@ -237,7 +239,9 @@ pub fn expr_trailing_brace(mut expr: &ast::Expr) -> Option<TrailingBrace<'_>> {
             | Yeet(None)
             | UnsafeBinderCast(..)
             | Err(_)
-            | Dummy => break None,
+            | Dummy => {
+                break None;
+            }
         }
     }
 }
diff --git a/compiler/rustc_ast/src/visit.rs b/compiler/rustc_ast/src/visit.rs
index a72e834e17a..ce8d6df75af 100644
--- a/compiler/rustc_ast/src/visit.rs
+++ b/compiler/rustc_ast/src/visit.rs
@@ -1269,8 +1269,8 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expression: &'a Expr) -> V
             try_visit!(visitor.visit_ty(container));
             walk_list!(visitor, visit_ident, fields.iter());
         }
-        ExprKind::Yield(optional_expression, _) => {
-            visit_opt!(visitor, visit_expr, optional_expression);
+        ExprKind::Yield(kind) => {
+            visit_opt!(visitor, visit_expr, kind.expr());
         }
         ExprKind::Try(subexpression) => try_visit!(visitor.visit_expr(subexpression)),
         ExprKind::TryBlock(body) => try_visit!(visitor.visit_block(body)),
diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs
index 19a08f92ce7..7f3d060bb8a 100644
--- a/compiler/rustc_ast_lowering/src/expr.rs
+++ b/compiler/rustc_ast_lowering/src/expr.rs
@@ -351,7 +351,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                         rest,
                     )
                 }
-                ExprKind::Yield(opt_expr, _) => self.lower_expr_yield(e.span, opt_expr.as_deref()),
+                ExprKind::Yield(kind) => self.lower_expr_yield(e.span, kind.expr().map(|x| &**x)),
                 ExprKind::Err(guar) => hir::ExprKind::Err(*guar),
 
                 ExprKind::UnsafeBinderCast(kind, expr, ty) => hir::ExprKind::UnsafeBinderCast(
diff --git a/compiler/rustc_ast_lowering/src/format.rs b/compiler/rustc_ast_lowering/src/format.rs
index 2bbf957feeb..faa47274f96 100644
--- a/compiler/rustc_ast_lowering/src/format.rs
+++ b/compiler/rustc_ast_lowering/src/format.rs
@@ -642,7 +642,7 @@ fn may_contain_yield_point(e: &ast::Expr) -> bool {
         type Result = ControlFlow<()>;
 
         fn visit_expr(&mut self, e: &ast::Expr) -> ControlFlow<()> {
-            if let ast::ExprKind::Await(_, _) | ast::ExprKind::Yield(_, _) = e.kind {
+            if let ast::ExprKind::Await(_, _) | ast::ExprKind::Yield(_) = e.kind {
                 ControlFlow::Break(())
             } else {
                 visit::walk_expr(self, e)
diff --git a/compiler/rustc_ast_pretty/src/pprust/state/expr.rs b/compiler/rustc_ast_pretty/src/pprust/state/expr.rs
index caba5f2721a..7d9dc89bd75 100644
--- a/compiler/rustc_ast_pretty/src/pprust/state/expr.rs
+++ b/compiler/rustc_ast_pretty/src/pprust/state/expr.rs
@@ -761,7 +761,7 @@ impl<'a> State<'a> {
                 self.print_expr(e, FixupContext::default());
                 self.pclose();
             }
-            ast::ExprKind::Yield(e, YieldKind::Prefix) => {
+            ast::ExprKind::Yield(YieldKind::Prefix(e)) => {
                 self.word("yield");
 
                 if let Some(expr) = e {
@@ -773,9 +773,7 @@ impl<'a> State<'a> {
                     );
                 }
             }
-            ast::ExprKind::Yield(e, YieldKind::Postfix) => {
-                // It's not possible to have a postfix yield with no expression.
-                let e = e.as_ref().unwrap();
+            ast::ExprKind::Yield(YieldKind::Postfix(e)) => {
                 self.print_expr_cond_paren(
                     e,
                     e.precedence() < ExprPrecedence::Unambiguous,
diff --git a/compiler/rustc_builtin_macros/src/assert/context.rs b/compiler/rustc_builtin_macros/src/assert/context.rs
index 54c3cebedfc..a949ab94f3a 100644
--- a/compiler/rustc_builtin_macros/src/assert/context.rs
+++ b/compiler/rustc_builtin_macros/src/assert/context.rs
@@ -323,7 +323,7 @@ impl<'cx, 'a> Context<'cx, 'a> {
             | ExprKind::While(_, _, _)
             | ExprKind::Yeet(_)
             | ExprKind::Become(_)
-            | ExprKind::Yield(_, _)
+            | ExprKind::Yield(_)
             | ExprKind::UnsafeBinderCast(..) => {}
         }
     }
diff --git a/compiler/rustc_parse/src/parser/expr.rs b/compiler/rustc_parse/src/parser/expr.rs
index 15a625314b4..fc9a511d56a 100644
--- a/compiler/rustc_parse/src/parser/expr.rs
+++ b/compiler/rustc_parse/src/parser/expr.rs
@@ -1315,8 +1315,9 @@ impl<'a> Parser<'a> {
         if self.eat_keyword(exp!(Yield)) {
             let yield_span = self.prev_token.span;
             self.psess.gated_spans.gate(sym::yield_expr, yield_span);
-            return Ok(self
-                .mk_expr(lo.to(yield_span), ExprKind::Yield(Some(self_arg), YieldKind::Postfix)));
+            return Ok(
+                self.mk_expr(lo.to(yield_span), ExprKind::Yield(YieldKind::Postfix(self_arg)))
+            );
         }
 
         let fn_span_lo = self.token.span;
@@ -1893,7 +1894,7 @@ impl<'a> Parser<'a> {
     /// Parse `"yield" expr?`.
     fn parse_expr_yield(&mut self) -> PResult<'a, P<Expr>> {
         let lo = self.prev_token.span;
-        let kind = ExprKind::Yield(self.parse_expr_opt()?, YieldKind::Prefix);
+        let kind = ExprKind::Yield(YieldKind::Prefix(self.parse_expr_opt()?));
         let span = lo.to(self.prev_token.span);
         self.psess.gated_spans.gate(sym::yield_expr, span);
         let expr = self.mk_expr(span, kind);
@@ -4047,7 +4048,7 @@ impl MutVisitor for CondChecker<'_> {
             | ExprKind::MacCall(_)
             | ExprKind::Struct(_)
             | ExprKind::Repeat(_, _)
-            | ExprKind::Yield(_, _)
+            | ExprKind::Yield(_)
             | ExprKind::Yeet(_)
             | ExprKind::Become(_)
             | ExprKind::IncludedBytes(_)
diff --git a/src/tools/clippy/clippy_lints/src/suspicious_operation_groupings.rs b/src/tools/clippy/clippy_lints/src/suspicious_operation_groupings.rs
index 206912d8de4..0d809c17989 100644
--- a/src/tools/clippy/clippy_lints/src/suspicious_operation_groupings.rs
+++ b/src/tools/clippy/clippy_lints/src/suspicious_operation_groupings.rs
@@ -528,7 +528,7 @@ fn ident_difference_expr_with_base_location(
         &strip_non_ident_wrappers(left).kind,
         &strip_non_ident_wrappers(right).kind,
     ) {
-        (Yield(_, _), Yield(_, _))
+        (Yield(_), Yield(_))
         | (Try(_), Try(_))
         | (Paren(_), Paren(_))
         | (Repeat(_, _), Repeat(_, _))
diff --git a/src/tools/clippy/clippy_utils/src/ast_utils/mod.rs b/src/tools/clippy/clippy_utils/src/ast_utils/mod.rs
index deda6030831..54261079fca 100644
--- a/src/tools/clippy/clippy_utils/src/ast_utils/mod.rs
+++ b/src/tools/clippy/clippy_utils/src/ast_utils/mod.rs
@@ -201,7 +201,7 @@ pub fn eq_expr(l: &Expr, r: &Expr) -> bool {
         (Loop(lt, ll, _), Loop(rt, rl, _)) => eq_label(ll.as_ref(), rl.as_ref()) && eq_block(lt, rt),
         (Block(lb, ll), Block(rb, rl)) => eq_label(ll.as_ref(), rl.as_ref()) && eq_block(lb, rb),
         (TryBlock(l), TryBlock(r)) => eq_block(l, r),
-        (Yield(l, lk), Yield(r, rk)) => eq_expr_opt(l.as_ref(), r.as_ref()) && lk == rk,
+        (Yield(l), Yield(r)) => eq_expr_opt(l.expr(), r.expr()) && l.same_kind(r),
         (Ret(l), Ret(r)) => eq_expr_opt(l.as_ref(), r.as_ref()),
         (Break(ll, le), Break(rl, re)) => eq_label(ll.as_ref(), rl.as_ref()) && eq_expr_opt(le.as_ref(), re.as_ref()),
         (Continue(ll), Continue(rl)) => eq_label(ll.as_ref(), rl.as_ref()),
diff --git a/src/tools/rustfmt/src/chains.rs b/src/tools/rustfmt/src/chains.rs
index fabb4400553..034ecde068a 100644
--- a/src/tools/rustfmt/src/chains.rs
+++ b/src/tools/rustfmt/src/chains.rs
@@ -259,7 +259,7 @@ impl ChainItemKind {
                 let span = mk_sp(nested.span.hi(), expr.span.hi());
                 (ChainItemKind::Await, span)
             }
-            ast::ExprKind::Yield(Some(ref nested), ast::YieldKind::Postfix) => {
+            ast::ExprKind::Yield(ast::YieldKind::Postfix(ref nested)) => {
                 let span = mk_sp(nested.span.hi(), expr.span.hi());
                 (ChainItemKind::Yield, span)
             }
@@ -516,7 +516,7 @@ impl Chain {
             ast::ExprKind::Field(ref subexpr, _)
             | ast::ExprKind::Try(ref subexpr)
             | ast::ExprKind::Await(ref subexpr, _)
-            | ast::ExprKind::Yield(Some(ref subexpr), ast::YieldKind::Postfix) => Some(SubExpr {
+            | ast::ExprKind::Yield(ast::YieldKind::Postfix(ref subexpr)) => Some(SubExpr {
                 expr: Self::convert_try(subexpr, context),
                 is_method_call_receiver: false,
             }),
diff --git a/src/tools/rustfmt/src/expr.rs b/src/tools/rustfmt/src/expr.rs
index 92c1ffa6076..e866f13efc7 100644
--- a/src/tools/rustfmt/src/expr.rs
+++ b/src/tools/rustfmt/src/expr.rs
@@ -221,7 +221,7 @@ pub(crate) fn format_expr(
                 Ok(format!("break{id_str}"))
             }
         }
-        ast::ExprKind::Yield(ref opt_expr, ast::YieldKind::Prefix) => {
+        ast::ExprKind::Yield(ast::YieldKind::Prefix(ref opt_expr)) => {
             if let Some(ref expr) = *opt_expr {
                 rewrite_unary_prefix(context, "yield ", &**expr, shape)
             } else {
@@ -244,7 +244,7 @@ pub(crate) fn format_expr(
         | ast::ExprKind::Field(..)
         | ast::ExprKind::MethodCall(..)
         | ast::ExprKind::Await(_, _)
-        | ast::ExprKind::Yield(_, ast::YieldKind::Postfix) => rewrite_chain(expr, context, shape),
+        | ast::ExprKind::Yield(ast::YieldKind::Postfix(_)) => rewrite_chain(expr, context, shape),
         ast::ExprKind::MacCall(ref mac) => {
             rewrite_macro(mac, None, context, shape, MacroPosition::Expression).or_else(|_| {
                 wrap_str(
diff --git a/src/tools/rustfmt/src/utils.rs b/src/tools/rustfmt/src/utils.rs
index 1811752c3c4..fcd475b1784 100644
--- a/src/tools/rustfmt/src/utils.rs
+++ b/src/tools/rustfmt/src/utils.rs
@@ -485,7 +485,7 @@ pub(crate) fn is_block_expr(context: &RewriteContext<'_>, expr: &ast::Expr, repr
         | ast::ExprKind::Index(_, ref expr, _)
         | ast::ExprKind::Unary(_, ref expr)
         | ast::ExprKind::Try(ref expr)
-        | ast::ExprKind::Yield(Some(ref expr), YieldKind::Prefix) => {
+        | ast::ExprKind::Yield(YieldKind::Prefix(Some(ref expr))) => {
             is_block_expr(context, expr, repr)
         }
         ast::ExprKind::Closure(ref closure) => is_block_expr(context, &closure.body, repr),
@@ -517,7 +517,7 @@ pub(crate) fn is_block_expr(context: &RewriteContext<'_>, expr: &ast::Expr, repr
         | ast::ExprKind::Tup(..)
         | ast::ExprKind::Use(..)
         | ast::ExprKind::Type(..)
-        | ast::ExprKind::Yield(_, _)
+        | ast::ExprKind::Yield(..)
         | ast::ExprKind::Underscore => false,
     }
 }