Update conditionalAssignment rule to support more complex lvalues

nicklockwood · Mar 18, 2024 · fdddd71 · fdddd71
1 parent 9ce956c
commit fdddd71
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 13 deletions.
diff --git a/Sources/Rules.swift b/Sources/Rules.swift
@@ -7107,11 +7107,10 @@ public struct _FormatRules {
                 startOfFirstBranch = startOfNestedBranch
             }
 
-            // Check if the first branch starts with the pattern `identifier =`.
-            guard let firstIdentifierIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: startOfFirstBranch),
-                  let identifier = formatter.token(at: firstIdentifierIndex),
-                  identifier.isIdentifier,
-                  let equalsIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: firstIdentifierIndex),
+            // Check if the first branch starts with the pattern `lvalue =`.
+            guard let firstTokenIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: startOfFirstBranch),
+                  let lvalueRange = formatter.parseExpressionRange(startingAt: firstTokenIndex),
+                  let equalsIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: lvalueRange.upperBound),
                   formatter.tokens[equalsIndex] == .operator("=", .infix)
             else { return }
 
@@ -7142,7 +7141,7 @@ public struct _FormatRules {
 
             // Whether or not the given conditional branch body qualifies as a single statement
             // that assigns a value to `identifier`. This is either:
-            //  1. a single assignment to `identifier =`
+            //  1. a single assignment to `lvalue =`
             //  2. a single `if` or `switch` statement where each of the branches also qualify,
             //     and the statement is exhaustive.
             func isExhaustiveSingleStatementAssignment(_ branch: Formatter.ConditionalBranch) -> Bool {
@@ -7169,9 +7168,10 @@ public struct _FormatRules {
                         && isExhaustive
                 }
 
-                // Otherwise we expect this to be of the pattern `identifier = (statement)`
-                else if formatter.tokens[firstTokenIndex] == identifier,
-                        let equalsIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: firstTokenIndex),
+                // Otherwise we expect this to be of the pattern `lvalue = (statement)`
+                else if let firstExpressionRange = formatter.parseExpressionRange(startingAt: firstTokenIndex),
+                        formatter.tokens[firstExpressionRange] == formatter.tokens[lvalueRange],
+                        let equalsIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: firstExpressionRange.upperBound),
                         formatter.tokens[equalsIndex] == .operator("=", .infix),
                         let valueStartIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: equalsIndex)
                 {
@@ -7223,7 +7223,8 @@ public struct _FormatRules {
             func removeAssignmentFromAllBranches() {
                 formatter.forEachRecursiveConditionalBranch(in: conditionalBranches) { branch in
                     guard let firstTokenIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: branch.startOfBranch),
-                          let equalsIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: firstTokenIndex),
+                          let firstExpressionRange = formatter.parseExpressionRange(startingAt: firstTokenIndex),
+                          let equalsIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: firstExpressionRange.upperBound),
                           formatter.tokens[equalsIndex] == .operator("=", .infix),
                           let valueStartIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: equalsIndex)
                     else { return }
@@ -7241,7 +7242,8 @@ public struct _FormatRules {
                let propertyIdentifierIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: introducerIndex),
                let propertyIdentifier = formatter.token(at: propertyIdentifierIndex),
                propertyIdentifier.isIdentifier,
-               propertyIdentifier == identifier,
+               formatter.tokens[lvalueRange.lowerBound] == propertyIdentifier,
+               lvalueRange.count == 1,
                let colonIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: propertyIdentifierIndex),
                formatter.tokens[colonIndex] == .delimiter(":"),
                let startOfTypeIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: colonIndex),
@@ -7305,12 +7307,13 @@ public struct _FormatRules {
                    (startOfFirstParentBranch ... endOfLastParentBranch).contains(startOfConditional)
                 { return }
 
+                let lvalueTokens = formatter.tokens[lvalueRange]
+
                 // Now we can remove the `identifier =` from each branch,
                 // and instead add it before the if / switch expression.
                 removeAssignmentFromAllBranches()
 
-                let identifierEqualsTokens: [Token] = [
-                    identifier,
+                let identifierEqualsTokens = lvalueTokens + [
                     .space(" "),
                     .operator("=", .infix),
                     .space(" "),

diff --git a/Tests/RulesTests+Syntax.swift b/Tests/RulesTests+Syntax.swift
@@ -4354,6 +4354,30 @@ class SyntaxTests: RulesTests {
         testFormatting(for: input, [output], rules: [FormatRules.conditionalAssignment, FormatRules.wrapMultilineConditionalAssignment, FormatRules.indent], options: options)
     }
 
+    func testConvertsSwitchStatementWithComplexLValueNotFollowingPropertyDefinition() {
+        let input = """
+        switch condition {
+        case true:
+            property?.foo!.bar["baaz"] = Foo("foo")
+        case false:
+            property?.foo!.bar["baaz"] = Foo("bar")
+        }
+        """
+
+        let output = """
+        property?.foo!.bar["baaz"] =
+            switch condition {
+            case true:
+                Foo("foo")
+            case false:
+                Foo("bar")
+            }
+        """
+
+        let options = FormatOptions(swiftVersion: "5.9")
+        testFormatting(for: input, [output], rules: [FormatRules.conditionalAssignment, FormatRules.wrapMultilineConditionalAssignment, FormatRules.indent], options: options)
+    }
+
     func testDoesntMergePropertyWithUnrelatedCondition() {
         let input = """
         let differentProperty: Foo

diff --git a/Tests/RulesTests+Wrapping.swift b/Tests/RulesTests+Wrapping.swift
@@ -5494,4 +5494,27 @@ class WrappingTests: RulesTests {
 
         testFormatting(for: input, [output], rules: [FormatRules.wrapMultilineConditionalAssignment, FormatRules.indent])
     }
+
+    func testWrapSwitchAssignmentWithComplexLValue() {
+        let input = """
+        property?.foo!.bar["baaz"] = switch condition {
+        case true:
+            Foo("foo")
+        case false:
+            Foo("bar")
+        }
+        """
+
+        let output = """
+        property?.foo!.bar["baaz"] =
+            switch condition {
+            case true:
+                Foo("foo")
+            case false:
+                Foo("bar")
+            }
+        """
+
+        testFormatting(for: input, [output], rules: [FormatRules.wrapMultilineConditionalAssignment, FormatRules.indent])
+    }
 }