Preserve yield position during pretty printing

4 files changed, 14 insertions, 5 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)),