Implement asymmetrical precedence for closures and jumps

2 files changed, 134 insertions, 64 deletions

diff --git a/compiler/rustc_ast_pretty/src/pprust/state/expr.rs b/compiler/rustc_ast_pretty/src/pprust/state/expr.rs
index c9a7e2aebd0..ee49246a4bb 100644
--- a/compiler/rustc_ast_pretty/src/pprust/state/expr.rs
+++ b/compiler/rustc_ast_pretty/src/pprust/state/expr.rs
@@ -7,8 +7,8 @@ use rustc_ast::util::classify;
 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, YieldKind, token,
+    self as ast, BinOpKind, BlockCheckMode, FormatAlignment, FormatArgPosition, FormatArgsPiece,
+    FormatCount, FormatDebugHex, FormatSign, FormatTrait, YieldKind, token,
 };
 
 use crate::pp::Breaks::Inconsistent;
@@ -214,13 +214,6 @@ impl<'a> State<'a> {
     }
 
     fn print_expr_call(&mut self, func: &ast::Expr, args: &[P<ast::Expr>], fixup: FixupContext) {
-        let needs_paren = match func.kind {
-            // In order to call a named field, needs parens: `(self.fun)()`
-            // But not for an unnamed field: `self.0()`
-            ast::ExprKind::Field(_, name) => !name.is_numeric(),
-            _ => func.precedence() < ExprPrecedence::Unambiguous,
-        };
-
         // Independent of parenthesization related to precedence, we must
         // parenthesize `func` if this is a statement context in which without
         // parentheses, a statement boundary would occur inside `func` or
@@ -237,8 +230,16 @@ impl<'a> State<'a> {
         // because the latter is valid syntax but with the incorrect meaning.
         // It's a match-expression followed by tuple-expression, not a function
         // call.
-        self.print_expr_cond_paren(func, needs_paren, fixup.leftmost_subexpression());
+        let func_fixup = fixup.leftmost_subexpression_with_operator(true);
+
+        let needs_paren = match func.kind {
+            // In order to call a named field, needs parens: `(self.fun)()`
+            // But not for an unnamed field: `self.0()`
+            ast::ExprKind::Field(_, name) => !name.is_numeric(),
+            _ => func_fixup.precedence(func) < ExprPrecedence::Unambiguous,
+        };
 
+        self.print_expr_cond_paren(func, needs_paren, func_fixup);
         self.print_call_post(args)
     }
 
@@ -281,9 +282,24 @@ impl<'a> State<'a> {
         rhs: &ast::Expr,
         fixup: FixupContext,
     ) {
+        let operator_can_begin_expr = match op {
+            | BinOpKind::Sub     // -x
+            | BinOpKind::Mul     // *x
+            | BinOpKind::And     // &&x
+            | BinOpKind::Or      // || x
+            | BinOpKind::BitAnd  // &x
+            | BinOpKind::BitOr   // |x| x
+            | BinOpKind::Shl     // <<T as Trait>::Type as Trait>::CONST
+            | BinOpKind::Lt      // <T as Trait>::CONST
+              => true,
+            _ => false,
+        };
+
+        let left_fixup = fixup.leftmost_subexpression_with_operator(operator_can_begin_expr);
+
         let binop_prec = op.precedence();
-        let left_prec = lhs.precedence();
-        let right_prec = rhs.precedence();
+        let left_prec = left_fixup.precedence(lhs);
+        let right_prec = fixup.precedence(rhs);
 
         let (mut left_needs_paren, right_needs_paren) = match op.fixity() {
             Fixity::Left => (left_prec < binop_prec, right_prec <= binop_prec),
@@ -312,18 +328,18 @@ impl<'a> State<'a> {
             _ => {}
         }
 
-        self.print_expr_cond_paren(lhs, left_needs_paren, fixup.leftmost_subexpression());
+        self.print_expr_cond_paren(lhs, left_needs_paren, left_fixup);
         self.space();
         self.word_space(op.as_str());
-        self.print_expr_cond_paren(rhs, right_needs_paren, fixup.subsequent_subexpression());
+        self.print_expr_cond_paren(rhs, right_needs_paren, fixup.rightmost_subexpression());
     }
 
     fn print_expr_unary(&mut self, op: ast::UnOp, expr: &ast::Expr, fixup: FixupContext) {
         self.word(op.as_str());
         self.print_expr_cond_paren(
             expr,
-            expr.precedence() < ExprPrecedence::Prefix,
-            fixup.subsequent_subexpression(),
+            fixup.precedence(expr) < ExprPrecedence::Prefix,
+            fixup.rightmost_subexpression(),
         );
     }
 
@@ -344,8 +360,8 @@ impl<'a> State<'a> {
         }
         self.print_expr_cond_paren(
             expr,
-            expr.precedence() < ExprPrecedence::Prefix,
-            fixup.subsequent_subexpression(),
+            fixup.precedence(expr) < ExprPrecedence::Prefix,
+            fixup.rightmost_subexpression(),
         );
     }
 
@@ -590,8 +606,8 @@ impl<'a> State<'a> {
                 self.word_space("=");
                 self.print_expr_cond_paren(
                     rhs,
-                    rhs.precedence() < ExprPrecedence::Assign,
-                    fixup.subsequent_subexpression(),
+                    fixup.precedence(rhs) < ExprPrecedence::Assign,
+                    fixup.rightmost_subexpression(),
                 );
             }
             ast::ExprKind::AssignOp(op, lhs, rhs) => {
@@ -604,8 +620,8 @@ impl<'a> State<'a> {
                 self.word_space(op.node.as_str());
                 self.print_expr_cond_paren(
                     rhs,
-                    rhs.precedence() < ExprPrecedence::Assign,
-                    fixup.subsequent_subexpression(),
+                    fixup.precedence(rhs) < ExprPrecedence::Assign,
+                    fixup.rightmost_subexpression(),
                 );
             }
             ast::ExprKind::Field(expr, ident) => {
@@ -618,10 +634,11 @@ impl<'a> State<'a> {
                 self.print_ident(*ident);
             }
             ast::ExprKind::Index(expr, index, _) => {
+                let expr_fixup = fixup.leftmost_subexpression_with_operator(true);
                 self.print_expr_cond_paren(
                     expr,
-                    expr.precedence() < ExprPrecedence::Unambiguous,
-                    fixup.leftmost_subexpression(),
+                    expr_fixup.precedence(expr) < ExprPrecedence::Unambiguous,
+                    expr_fixup,
                 );
                 self.word("[");
                 self.print_expr(index, FixupContext::default());
@@ -634,10 +651,11 @@ impl<'a> State<'a> {
                 // a "normal" binop gets parenthesized. (`LOr` is the lowest-precedence binop.)
                 let fake_prec = ExprPrecedence::LOr;
                 if let Some(e) = start {
+                    let start_fixup = fixup.leftmost_subexpression_with_operator(true);
                     self.print_expr_cond_paren(
                         e,
-                        e.precedence() < fake_prec,
-                        fixup.leftmost_subexpression(),
+                        start_fixup.precedence(e) < fake_prec,
+                        start_fixup,
                     );
                 }
                 match limits {
@@ -647,8 +665,8 @@ impl<'a> State<'a> {
                 if let Some(e) = end {
                     self.print_expr_cond_paren(
                         e,
-                        e.precedence() < fake_prec,
-                        fixup.subsequent_subexpression(),
+                        fixup.precedence(e) < fake_prec,
+                        fixup.rightmost_subexpression(),
                     );
                 }
             }
@@ -665,11 +683,10 @@ impl<'a> State<'a> {
                     self.space();
                     self.print_expr_cond_paren(
                         expr,
-                        // Parenthesize if required by precedence, or in the
-                        // case of `break 'inner: loop { break 'inner 1 } + 1`
-                        expr.precedence() < ExprPrecedence::Jump
-                            || (opt_label.is_none() && classify::leading_labeled_expr(expr)),
-                        fixup.subsequent_subexpression(),
+                        // Parenthesize `break 'inner: loop { break 'inner 1 } + 1`
+                        //                     ^---------------------------------^
+                        opt_label.is_none() && classify::leading_labeled_expr(expr),
+                        fixup.rightmost_subexpression(),
                     );
                 }
             }
@@ -684,11 +701,7 @@ impl<'a> State<'a> {
                 self.word("return");
                 if let Some(expr) = result {
                     self.word(" ");
-                    self.print_expr_cond_paren(
-                        expr,
-                        expr.precedence() < ExprPrecedence::Jump,
-                        fixup.subsequent_subexpression(),
-                    );
+                    self.print_expr(expr, fixup.rightmost_subexpression());
                 }
             }
             ast::ExprKind::Yeet(result) => {
@@ -697,21 +710,13 @@ impl<'a> State<'a> {
                 self.word("yeet");
                 if let Some(expr) = result {
                     self.word(" ");
-                    self.print_expr_cond_paren(
-                        expr,
-                        expr.precedence() < ExprPrecedence::Jump,
-                        fixup.subsequent_subexpression(),
-                    );
+                    self.print_expr(expr, fixup.rightmost_subexpression());
                 }
             }
             ast::ExprKind::Become(result) => {
                 self.word("become");
                 self.word(" ");
-                self.print_expr_cond_paren(
-                    result,
-                    result.precedence() < ExprPrecedence::Jump,
-                    fixup.subsequent_subexpression(),
-                );
+                self.print_expr(result, fixup.rightmost_subexpression());
             }
             ast::ExprKind::InlineAsm(a) => {
                 // FIXME: Print `builtin # asm` once macro `asm` uses `builtin_syntax`.
@@ -761,11 +766,7 @@ impl<'a> State<'a> {
 
                 if let Some(expr) = e {
                     self.space();
-                    self.print_expr_cond_paren(
-                        expr,
-                        expr.precedence() < ExprPrecedence::Jump,
-                        fixup.subsequent_subexpression(),
-                    );
+                    self.print_expr(expr, fixup.rightmost_subexpression());
                 }
             }
             ast::ExprKind::Yield(YieldKind::Postfix(e)) => {
diff --git a/compiler/rustc_ast_pretty/src/pprust/state/fixup.rs b/compiler/rustc_ast_pretty/src/pprust/state/fixup.rs
index 3ef21f5cb29..eb5ac8b78a8 100644
--- a/compiler/rustc_ast_pretty/src/pprust/state/fixup.rs
+++ b/compiler/rustc_ast_pretty/src/pprust/state/fixup.rs
@@ -1,5 +1,6 @@
-use rustc_ast::Expr;
-use rustc_ast::util::{classify, parser};
+use rustc_ast::util::classify;
+use rustc_ast::util::parser::{self, ExprPrecedence};
+use rustc_ast::{Expr, ExprKind, YieldKind};
 
 // The default amount of fixing is minimal fixing, so all fixups are set to `false` by `Default`.
 // Fixups should be turned on in a targeted fashion where needed.
@@ -93,6 +94,24 @@ pub(crate) struct FixupContext {
     /// }
     /// ```
     parenthesize_exterior_struct_lit: bool,
+
+    /// This is the difference between:
+    ///
+    /// ```ignore (illustrative)
+    /// let _ = (return) - 1;  // without paren, this would return -1
+    ///
+    /// let _ = return + 1;  // no paren because '+' cannot begin expr
+    /// ```
+    next_operator_can_begin_expr: bool,
+
+    /// This is the difference between:
+    ///
+    /// ```ignore (illustrative)
+    /// let _ = 1 + return 1;  // no parens if rightmost subexpression
+    ///
+    /// let _ = 1 + (return 1) + 1;  // needs parens
+    /// ```
+    next_operator_can_continue_expr: bool,
 }
 
 impl FixupContext {
@@ -134,6 +153,8 @@ impl FixupContext {
             match_arm: false,
             leftmost_subexpression_in_match_arm: self.match_arm
                 || self.leftmost_subexpression_in_match_arm,
+            next_operator_can_begin_expr: false,
+            next_operator_can_continue_expr: true,
             ..self
         }
     }
@@ -148,19 +169,34 @@ impl FixupContext {
             leftmost_subexpression_in_stmt: false,
             match_arm: self.match_arm || self.leftmost_subexpression_in_match_arm,
             leftmost_subexpression_in_match_arm: false,
+            next_operator_can_begin_expr: false,
+            next_operator_can_continue_expr: true,
             ..self
         }
     }
 
-    /// Transform this fixup into the one that should apply when printing any
-    /// subexpression that is neither a leftmost subexpression nor surrounded in
-    /// delimiters.
+    /// Transform this fixup into the one that should apply when printing a
+    /// leftmost subexpression followed by punctuation that is legal as the
+    /// first token of an expression.
+    pub(crate) fn leftmost_subexpression_with_operator(
+        self,
+        next_operator_can_begin_expr: bool,
+    ) -> Self {
+        FixupContext { next_operator_can_begin_expr, ..self.leftmost_subexpression() }
+    }
+
+    /// Transform this fixup into the one that should apply when printing the
+    /// rightmost subexpression of the current expression.
+    ///
+    /// The rightmost subexpression is any subexpression that has a different
+    /// first token than the current expression, but has the same last token.
+    ///
+    /// For example in `$a + $b` and `-$b`, the subexpression `$b` is a
+    /// rightmost subexpression.
     ///
-    /// This is for any subexpression that has a different first token than the
-    /// current expression, and is not surrounded by a paren/bracket/brace. For
-    /// example the `$b` in `$a + $b` and `-$b`, but not the one in `[$b]` or
-    /// `$a.f($b)`.
-    pub(crate) fn subsequent_subexpression(self) -> Self {
+    /// Not every expression has a rightmost subexpression. For example neither
+    /// `[$b]` nor `$a.f($b)` have one.
+    pub(crate) fn rightmost_subexpression(self) -> Self {
         FixupContext {
             stmt: false,
             leftmost_subexpression_in_stmt: false,
@@ -193,6 +229,39 @@ impl FixupContext {
     ///     "let chain".
     pub(crate) fn needs_par_as_let_scrutinee(self, expr: &Expr) -> bool {
         self.parenthesize_exterior_struct_lit && parser::contains_exterior_struct_lit(expr)
-            || parser::needs_par_as_let_scrutinee(expr.precedence())
+            || parser::needs_par_as_let_scrutinee(self.precedence(expr))
+    }
+
+    /// Determines the effective precedence of a subexpression. Some expressions
+    /// have higher or lower precedence when adjacent to particular operators.
+    pub(crate) fn precedence(self, expr: &Expr) -> ExprPrecedence {
+        if self.next_operator_can_begin_expr {
+            // Decrease precedence of value-less jumps when followed by an
+            // operator that would otherwise get interpreted as beginning a
+            // value for the jump.
+            if let ExprKind::Break(..)
+            | ExprKind::Ret(..)
+            | ExprKind::Yeet(..)
+            | ExprKind::Yield(YieldKind::Prefix(..)) = expr.kind
+            {
+                return ExprPrecedence::Jump;
+            }
+        }
+
+        if !self.next_operator_can_continue_expr {
+            // Increase precedence of expressions that extend to the end of
+            // current statement or group.
+            if let ExprKind::Break(..)
+            | ExprKind::Closure(..)
+            | ExprKind::Ret(..)
+            | ExprKind::Yeet(..)
+            | ExprKind::Yield(YieldKind::Prefix(..))
+            | ExprKind::Range(None, ..) = expr.kind
+            {
+                return ExprPrecedence::Prefix;
+            }
+        }
+
+        expr.precedence()
     }
 }