From 4fb999d84abcf960a1dfd165d7b25d99cc746f24 Mon Sep 17 00:00:00 2001
From: Tony Allevato <allevato@google.com>
Date: Thu, 25 May 2023 16:11:26 -0700
Subject: [PATCH] Fix `try`/`await` expressions in `NoAssignmentInExpressions`.

This rule only checked whether the _immediate_ parent of an expression
like `x = y` was a `CodeBlockItem`. In cases like `try x = y`, the `try`
wraps the entire expression and the `CodeBlockItem` would be the parent
of that. So we look through any `TryExpr`s or `AwaitExpr`s we see when
searching for the `CodeBlockItem`.
---
 .../NoAssignmentInExpressions.swift           | 29 +++++++++++++++++--
 .../NoAssignmentInExpressionsTests.swift      | 19 ++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/Sources/SwiftFormatRules/NoAssignmentInExpressions.swift b/Sources/SwiftFormatRules/NoAssignmentInExpressions.swift
index e327b59d7..0ae58473e 100644
--- a/Sources/SwiftFormatRules/NoAssignmentInExpressions.swift
+++ b/Sources/SwiftFormatRules/NoAssignmentInExpressions.swift
@@ -27,7 +27,7 @@ public final class NoAssignmentInExpressions: SyntaxFormatRule {
   public override func visit(_ node: InfixOperatorExprSyntax) -> ExprSyntax {
     // Diagnose any assignment that isn't directly a child of a `CodeBlockItem` (which would be the
     // case if it was its own statement).
-    if isAssignmentExpression(node) && node.parent?.is(CodeBlockItemSyntax.self) == false {
+    if isAssignmentExpression(node) && !isStandaloneAssignmentStatement(node) {
       diagnose(.moveAssignmentToOwnStatement, on: node)
     }
     return ExprSyntax(node)
@@ -59,7 +59,8 @@ public final class NoAssignmentInExpressions: SyntaxFormatRule {
             item: .expr(ExprSyntax(assignmentExpr)),
             semicolon: nil
           )
-          .with(\.leadingTrivia, 
+          .with(
+            \.leadingTrivia,
             (returnStmt.leadingTrivia) + (assignmentExpr.leadingTrivia))
           .with(\.trailingTrivia, []))
         newItems.append(
@@ -106,6 +107,30 @@ public final class NoAssignmentInExpressions: SyntaxFormatRule {
     return context.operatorTable.infixOperator(named: binaryOp.operatorToken.text)?.precedenceGroup
       == "AssignmentPrecedence"
   }
+
+  /// Returns a value indicating whether the given node is a standalone assignment statement.
+  ///
+  /// This function considers try/await expressions and automatically walks up through them as
+  /// needed. This is because `try f().x = y` should still be a standalone assignment for our
+  /// purposes, even though a `TryExpr` will wrap the `InfixOperatorExpr` and thus would not be
+  /// considered a standalone assignment if we only checked the infix expression for a
+  /// `CodeBlockItem` parent.
+  private func isStandaloneAssignmentStatement(_ node: InfixOperatorExprSyntax) -> Bool {
+    var node = Syntax(node)
+    while
+      let parent = node.parent,
+      parent.is(TryExprSyntax.self) || parent.is(AwaitExprSyntax.self)
+    {
+      node = parent
+    }
+
+    guard let parent = node.parent else {
+      // This shouldn't happen under normal circumstances (i.e., unless the expression is detached
+      // from the rest of a tree). In that case, we may as well consider it to be "standalone".
+      return true
+    }
+    return parent.is(CodeBlockItemSyntax.self)
+  }
 }
 
 extension Finding.Message {
diff --git a/Tests/SwiftFormatRulesTests/NoAssignmentInExpressionsTests.swift b/Tests/SwiftFormatRulesTests/NoAssignmentInExpressionsTests.swift
index 238c2f92f..dfeeefbf8 100644
--- a/Tests/SwiftFormatRulesTests/NoAssignmentInExpressionsTests.swift
+++ b/Tests/SwiftFormatRulesTests/NoAssignmentInExpressionsTests.swift
@@ -161,4 +161,23 @@ final class NoAssignmentInExpressionsTests: LintOrFormatRuleTestCase {
     )
     XCTAssertDiagnosed(.moveAssignmentToOwnStatement, line: 2, column: 29)
   }
+
+  func testTryAndAwaitAssignmentExpressionsAreUnchanged() {
+    XCTAssertFormatting(
+      NoAssignmentInExpressions.self,
+      input: """
+        func foo() {
+          try a.b = c
+          await a.b = c
+        }
+        """,
+      expected: """
+        func foo() {
+          try a.b = c
+          await a.b = c
+        }
+        """
+    )
+    XCTAssertNotDiagnosed(.moveAssignmentToOwnStatement)
+  }
 }