Preserve yield position during pretty printing

12 files changed, 55 insertions, 18 deletions

diff --git a/compiler/rustc_ast/src/ast.rs b/compiler/rustc_ast/src/ast.rs
index 1b831c454e6..9dcdd868343 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>>),
+    Yield(Option<P<Expr>>, YieldKind),
 
     /// A `do yeet` (aka `throw`/`fail`/`bail`/`raise`/whatever),
     /// with an optional value to be returned.
@@ -1903,6 +1903,15 @@ pub enum MatchKind {
     Postfix,
 }
 
+/// The kind of yield expression
+#[derive(Clone, Copy, Encodable, Decodable, Debug, PartialEq)]
+pub enum YieldKind {
+    /// yield expr { ... }
+    Prefix,
+    /// expr.yield { ... }
+    Postfix,
+}
+
 /// A literal in a meta item.
 #[derive(Clone, Encodable, Decodable, Debug, HashStable_Generic)]
 pub struct MetaItemLit {
diff --git a/compiler/rustc_ast/src/mut_visit.rs b/compiler/rustc_ast/src/mut_visit.rs
index 4a1636e6aec..5a9df6ffadf 100644
--- a/compiler/rustc_ast/src/mut_visit.rs
+++ b/compiler/rustc_ast/src/mut_visit.rs
@@ -1813,7 +1813,7 @@ 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) => {
+        ExprKind::Yield(expr, _) => {
             visit_opt(expr, |expr| vis.visit_expr(expr));
         }
         ExprKind::Try(expr) => vis.visit_expr(expr),
diff --git a/compiler/rustc_ast/src/util/classify.rs b/compiler/rustc_ast/src/util/classify.rs
index e43d78f6e72..116847af9f4 100644
--- a/compiler/rustc_ast/src/util/classify.rs
+++ b/compiler/rustc_ast/src/util/classify.rs
@@ -182,7 +182,7 @@ pub fn expr_trailing_brace(mut expr: &ast::Expr) -> Option<TrailingBrace<'_>> {
             | Range(_, Some(e), _)
             | Ret(Some(e))
             | Unary(_, e)
-            | Yield(Some(e))
+            | Yield(Some(e), _)
             | Yeet(Some(e))
             | Become(e) => {
                 expr = e;
@@ -217,7 +217,7 @@ pub fn expr_trailing_brace(mut expr: &ast::Expr) -> Option<TrailingBrace<'_>> {
             Break(_, None)
             | Range(_, None, _)
             | Ret(None)
-            | Yield(None)
+            | Yield(None, _)
             | Array(_)
             | Call(_, _)
             | MethodCall(_)
diff --git a/compiler/rustc_ast/src/visit.rs b/compiler/rustc_ast/src/visit.rs
index cfcb0e23cb5..a72e834e17a 100644
--- a/compiler/rustc_ast/src/visit.rs
+++ b/compiler/rustc_ast/src/visit.rs
@@ -1269,7 +1269,7 @@ 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) => {
+        ExprKind::Yield(optional_expression, _) => {
             visit_opt!(visitor, visit_expr, optional_expression);
         }
         ExprKind::Try(subexpression) => try_visit!(visitor.visit_expr(subexpression)),
diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs
index 5bb6704dde4..19a08f92ce7 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(opt_expr, _) => self.lower_expr_yield(e.span, opt_expr.as_deref()),
                 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 faa47274f96..2bbf957feeb 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 e3c41f117ab..9e53ed41c46 100644
--- a/compiler/rustc_ast_pretty/src/pprust/state/expr.rs
+++ b/compiler/rustc_ast_pretty/src/pprust/state/expr.rs
@@ -8,7 +8,7 @@ use rustc_ast::util::literal::escape_byte_str_symbol;
 use rustc_ast::util::parser::{self, ExprPrecedence, Fixity};
 use rustc_ast::{
     self as ast, BlockCheckMode, FormatAlignment, FormatArgPosition, FormatArgsPiece, FormatCount,
-    FormatDebugHex, FormatSign, FormatTrait, token,
+    FormatDebugHex, FormatSign, FormatTrait, YieldKind, token,
 };
 
 use crate::pp::Breaks::Inconsistent;
@@ -761,7 +761,7 @@ impl<'a> State<'a> {
                 self.print_expr(e, FixupContext::default());
                 self.pclose();
             }
-            ast::ExprKind::Yield(e) => {
+            ast::ExprKind::Yield(e, YieldKind::Prefix) => {
                 self.word("yield");
 
                 if let Some(expr) = e {
@@ -773,6 +773,16 @@ 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();
+                self.print_expr_cond_paren(
+                    e,
+                    e.precedence() < ExprPrecedence::Unambiguous,
+                    fixup.leftmost_subexpression_with_dot(),
+                );
+                self.word(".yield");
+            }
             ast::ExprKind::Try(e) => {
                 self.print_expr_cond_paren(
                     e,
diff --git a/compiler/rustc_builtin_macros/src/assert/context.rs b/compiler/rustc_builtin_macros/src/assert/context.rs
index a949ab94f3a..54c3cebedfc 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 1df6283af26..cb04bbc240e 100644
--- a/compiler/rustc_parse/src/parser/expr.rs
+++ b/compiler/rustc_parse/src/parser/expr.rs
@@ -17,6 +17,7 @@ use rustc_ast::{
     self as ast, AnonConst, Arm, AttrStyle, AttrVec, BinOp, BinOpKind, BlockCheckMode, CaptureBy,
     ClosureBinder, DUMMY_NODE_ID, Expr, ExprField, ExprKind, FnDecl, FnRetTy, Label, MacCall,
     MetaItemLit, Movability, Param, RangeLimits, StmtKind, Ty, TyKind, UnOp, UnsafeBinderCastKind,
+    YieldKind,
 };
 use rustc_ast_pretty::pprust;
 use rustc_data_structures::stack::ensure_sufficient_stack;
@@ -1314,7 +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(yield_span, ExprKind::Yield(Some(self_arg))));
+            return Ok(
+                self.mk_expr(yield_span, ExprKind::Yield(Some(self_arg), YieldKind::Postfix))
+            );
         }
 
         let fn_span_lo = self.token.span;
@@ -1891,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()?);
+        let kind = ExprKind::Yield(self.parse_expr_opt()?, YieldKind::Prefix);
         let span = lo.to(self.prev_token.span);
         self.psess.gated_spans.gate(sym::yield_expr, span);
         let expr = self.mk_expr(span, kind);
@@ -4045,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/rustfmt/src/utils.rs b/src/tools/rustfmt/src/utils.rs
index fe716c18638..bee39153229 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)) => is_block_expr(context, expr, repr),
+        | ast::ExprKind::Yield(Some(ref expr), _) => is_block_expr(context, expr, repr),
         ast::ExprKind::Closure(ref closure) => is_block_expr(context, &closure.body, repr),
         // This can only be a string lit
         ast::ExprKind::Lit(_) => {
@@ -515,7 +515,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(None)
+        | ast::ExprKind::Yield(None, _)
         | ast::ExprKind::Underscore => false,
     }
 }
diff --git a/tests/pretty/postfix-yield.rs b/tests/pretty/postfix-yield.rs
new file mode 100644
index 00000000000..f76e8142ae8
--- /dev/null
+++ b/tests/pretty/postfix-yield.rs
@@ -0,0 +1,15 @@
+// This demonstrates a proposed alternate or additional option of having yield in postfix position.
+//@ edition: 2024
+//@ pp-exact
+
+#![feature(gen_blocks, coroutines, coroutine_trait, yield_expr)]
+
+use std::ops::{Coroutine, CoroutineState};
+use std::pin::pin;
+
+fn main() {
+    let mut gn = gen { yield 1; 2.yield; (1 + 2).yield; };
+
+    let mut coro =
+        pin!(#[coroutine] |_: i32| { let x = 1.yield; (x + 2).yield; });
+}
diff --git a/tests/ui/coroutine/postfix-yield.rs b/tests/ui/coroutine/postfix-yield.rs
index 77b1c2a19d0..ff843138c8c 100644
--- a/tests/ui/coroutine/postfix-yield.rs
+++ b/tests/ui/coroutine/postfix-yield.rs
@@ -9,7 +9,7 @@ use std::ops::{Coroutine, CoroutineState};
 use std::pin::pin;
 
 fn main() {
-    // generators (i.e. yield doesn't return anything useful) 
+    // generators (i.e. yield doesn't return anything useful)
     let mut gn = gen {
         yield 1;
         2.yield;
@@ -23,8 +23,8 @@ fn main() {
     let mut coro = pin!(
         #[coroutine]
         |_: i32| {
-            let x = yield 1;
-            yield x + 2;
+            let x = 1.yield;
+            (x + 2).yield;
         }
     );