Teach rustfmt to handle postfix yield

This involved fixing the span when parsing .yield

Author: eholk

compiler/rustc_parse/src/parser/expr.rs

@@ -1315,9 +1315,8 @@ 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), YieldKind::Postfix))
-            );
+            return Ok(self
+                .mk_expr(lo.to(yield_span), ExprKind::Yield(Some(self_arg), YieldKind::Postfix)));
         }
 
         let fn_span_lo = self.token.span;

src/tools/rustfmt/src/chains.rs

@@ -192,6 +192,7 @@ enum ChainItemKind {
     StructField(symbol::Ident),
     TupleField(symbol::Ident, bool),
     Await,
+    Yield,
     Comment(String, CommentPosition),
 }
 
@@ -203,6 +204,7 @@ impl ChainItemKind {
             | ChainItemKind::StructField(..)
             | ChainItemKind::TupleField(..)
             | ChainItemKind::Await
+            | ChainItemKind::Yield
             | ChainItemKind::Comment(..) => false,
         }
     }
@@ -257,6 +259,10 @@ impl ChainItemKind {
                 let span = mk_sp(nested.span.hi(), expr.span.hi());
                 (ChainItemKind::Await, span)
             }
+            ast::ExprKind::Yield(Some(ref nested), ast::YieldKind::Postfix) => {
+                let span = mk_sp(nested.span.hi(), expr.span.hi());
+                (ChainItemKind::Yield, span)
+            }
             _ => {
                 return (
                     ChainItemKind::Parent {
@@ -306,6 +312,7 @@ impl Rewrite for ChainItem {
                 rewrite_ident(context, ident)
             ),
             ChainItemKind::Await => ".await".to_owned(),
+            ChainItemKind::Yield => ".yield".to_owned(),
             ChainItemKind::Comment(ref comment, _) => {
                 rewrite_comment(comment, false, shape, context.config)?
             }
@@ -508,7 +515,8 @@ impl Chain {
             }),
             ast::ExprKind::Field(ref subexpr, _)
             | ast::ExprKind::Try(ref subexpr)
-            | ast::ExprKind::Await(ref subexpr, _) => Some(SubExpr {
+            | ast::ExprKind::Await(ref subexpr, _)
+            | ast::ExprKind::Yield(Some(ref subexpr), ast::YieldKind::Postfix) => Some(SubExpr {
                 expr: Self::convert_try(subexpr, context),
                 is_method_call_receiver: false,
             }),

src/tools/rustfmt/src/utils.rs

@@ -4,7 +4,7 @@ use rustc_ast::ast::{
     self, Attribute, MetaItem, MetaItemInner, MetaItemKind, NodeId, Path, Visibility,
     VisibilityKind,
 };
-use rustc_ast::ptr;
+use rustc_ast::{YieldKind, ptr};
 use rustc_ast_pretty::pprust;
 use rustc_span::{BytePos, LocalExpnId, Span, Symbol, SyntaxContext, sym, symbol};
 use unicode_width::UnicodeWidthStr;
@@ -485,7 +485,9 @@ 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), YieldKind::Prefix) => {
+            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 +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(None, _)
+        | ast::ExprKind::Yield(_, _)
         | ast::ExprKind::Underscore => false,
     }
 }

src/tools/rustfmt/tests/source/postfix-yield.rs

@@ -7,6 +7,11 @@ use std::ops::{Coroutine, CoroutineState};
 use std::pin::pin;
 
 fn main() {
-    let mut coro =
-        pin!(#[coroutine] |_: i32| { let x = 1.yield; (x + 2).yield; });
+    let mut coro = pin!(
+        #[coroutine]
+        |_: i32| {
+            let x = 1.yield;
+            (x + 2).await;
+        }
+    );
 }