Implement wrapMultilineConditionalAssignment rule with WIP fix for in…

…dentation

Rules.md

@@ -80,6 +80,7 @@
 * [wrap](#wrap)
 * [wrapArguments](#wrapArguments)
 * [wrapAttributes](#wrapAttributes)
+* [wrapMultilineConditionalAssignment](#wrapMultilineConditionalAssignment)
 * [wrapMultilineStatementBraces](#wrapMultilineStatementBraces)
 * [wrapSingleLineComments](#wrapSingleLineComments)
 * [yodaConditions](#yodaConditions)
@@ -2691,6 +2692,28 @@ Option | Description
 </details>
 <br/>
 
+## wrapMultilineConditionalAssignment
+
+Wraps multiline conditional assignment expressions after the assignment operator
+
+<details>
+<summary>Examples</summary>
+
+- let planetLocation = if let star = planet.star {
+-     "The \(star.name) system"
+- } else {
+-     "Rogue planet"
+- }
++ let planetLocation =
++     if let star = planet.star {
++         "The \(star.name) system"
++     } else {
++         "Rogue planet"
++     }
+
+</details>
+<br/>
+
 ## wrapMultilineStatementBraces
 
 Wrap the opening brace of multiline statements.

Sources/Examples.swift

@@ -1737,4 +1737,18 @@ private struct Examples {
     + func foo(_ bar: Bar) { ... }
     ```
     """
+
+    let wrapMultilineConditionalAssignment = #"""
+    - let planetLocation = if let star = planet.star {
+    -     "The \(star.name) system"
+    - } else {
+    -     "Rogue planet"
+    - }
+    + let planetLocation =
+    +     if let star = planet.star {
+    +         "The \(star.name) system"
+    +     } else {
+    +         "Rogue planet"
+    +     }
+    """#
 }

Sources/Rules.swift

@@ -1959,7 +1959,24 @@ public struct _FormatRules {
                     } else if !formatter.options.xcodeIndentation || !isWrappedDeclaration() {
                         indent += formatter.linewrapIndent(at: i)
                     }
-                    linewrapStack[linewrapStack.count - 1] = true
+
+                    /// Whether or not this indent from a linewrap should be applied to following lines in the indentation scope.
+                    /// If true, doesn't modify the `linewrapStack`.
+                    var linewrapIndentAppliesToFollowingLines = false
+
+                    // If / switch expressions on the line following an = assignment
+                    // should be indented, including all of the branches.
+                    if let nextTokenIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: i),
+                       ["if", "switch"].contains(formatter.tokens[nextTokenIndex].string),
+                       !formatter.onSameLine(i, nextTokenIndex)
+                    {
+                        linewrapIndentAppliesToFollowingLines = true
+                    }
+
+                    if !linewrapIndentAppliesToFollowingLines {
+                        linewrapStack[linewrapStack.count - 1] = true
+                    }
+
                     indentStack.append(indent)
                     stringBodyIndentStack.append("")
                 }
@@ -7586,4 +7603,73 @@ public struct _FormatRules {
             }
         }
     }
+
+    public let wrapMultilineConditionalAssignment = FormatRule(
+        help: "Wraps multiline conditional assignment expressions after the assignment operator",
+        orderAfter: ["conditionalAssignment"]
+    ) { formatter in
+        formatter.forEach(.keyword) { introducerIndex, introducerToken in
+            guard ["let", "var"].contains(introducerToken.string),
+                  let identifierIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: introducerIndex),
+                  let identifier = formatter.token(at: identifierIndex),
+                  identifier.isIdentifier
+            else { return }
+
+            // Find the `=` index for this variable, if present
+            let assignmentIndex: Int
+            if let colonIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: identifierIndex),
+               formatter.tokens[colonIndex] == .delimiter(":"),
+               let startOfTypeIndex = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: colonIndex),
+               let typeRange = formatter.parseType(at: startOfTypeIndex)?.range,
+               let tokenAfterType = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: typeRange.upperBound),
+               formatter.tokens[tokenAfterType] == .operator("=", .infix)
+            {
+                assignmentIndex = tokenAfterType
+            }
+
+            else if let tokenAfterIdentifier = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: identifierIndex),
+                    formatter.tokens[tokenAfterIdentifier] == .operator("=", .infix)
+            {
+                assignmentIndex = tokenAfterIdentifier
+            }
+
+            else {
+                return
+            }
+
+            // Verify the RHS of the assignment is an if/switch expression
+            guard let startOfConditionalExpression = formatter.index(of: .nonSpaceOrCommentOrLinebreak, after: assignmentIndex),
+                  ["if", "switch"].contains(formatter.tokens[startOfConditionalExpression].string),
+                  let conditionalBranches = formatter.conditionalBranches(at: startOfConditionalExpression),
+                  let lastBranch = conditionalBranches.last
+            else { return }
+
+            // If the entire expression is on a single line, we leave the formatting as-is
+            guard !formatter.onSameLine(startOfConditionalExpression, lastBranch.endOfBranch) else {
+                return
+            }
+
+            // The `=` should be on the same line as the `let`/`var` introducer
+            if !formatter.onSameLine(introducerIndex, assignmentIndex),
+               formatter.last(.nonSpaceOrComment, before: assignmentIndex)?.isLinebreak == true,
+               let previousToken = formatter.index(of: .nonSpaceOrCommentOrLinebreak, before: assignmentIndex),
+               formatter.onSameLine(introducerIndex, previousToken)
+            {
+                // Move the assignment operator to follow the previous token.
+                // Also remove any trailing space after the previous position
+                // of the assignment operator.
+                if formatter.tokens[assignmentIndex + 1].isSpaceOrLinebreak {
+                    formatter.removeToken(at: assignmentIndex + 1)
+                }
+
+                formatter.removeToken(at: assignmentIndex)
+                formatter.insert([.space(" "), .operator("=", .infix)], at: previousToken + 1)
+            }
+
+            // And there should be a line break between the `=` and the `if` / `switch` keyword
+            else if !formatter.tokens[(assignmentIndex + 1) ..< startOfConditionalExpression].contains(where: \.isLinebreak) {
+                formatter.insertLinebreak(at: startOfConditionalExpression - 1)
+            }
+        }
+    }
 }

Tests/RulesTests+Indentation.swift

@@ -3759,4 +3759,126 @@ class IndentTests: RulesTests {
         let options = FormatOptions(closingParenOnSameLine: true)
         testFormatting(for: input, rule: FormatRules.indent, options: options)
     }
+
+    func testIndentIfExpressionAssignmentOnNextLine() {
+        let input = """
+        let foo =
+        if let bar {
+            bar
+        } else if let baaz {
+            baaz
+        } else if let quux {
+            if let foo {
+                foo
+            } else {
+                quux
+            }
+        }
+        """
+
+        let output = """
+        let foo =
+            if let bar {
+                bar
+            } else if let baaz {
+                baaz
+            } else if let quux {
+                if let foo {
+                    foo
+                } else {
+                    quux
+                }
+            }
+        """
+
+        testFormatting(for: input, output, rule: FormatRules.indent)
+    }
+
+    func testIndentIfExpressionAssignmentOnSameLine() {
+        let input = """
+        let foo = if let bar {
+            bar
+        } else if let baaz {
+            baaz
+        } else if let quux {
+            if let foo {
+                foo
+            } else {
+                quux
+            }
+        }
+        """
+
+        testFormatting(for: input, rule: FormatRules.indent, exclude: ["wrapMultilineConditionalAssignment"])
+    }
+
+    func testIndentSwitchExpressionAssignment() {
+        let input = """
+        let foo =
+        switch bar {
+        case true:
+            bar
+        case baaz:
+            baaz
+        }
+        """
+
+        let output = """
+        let foo =
+            switch bar {
+            case true:
+                bar
+            case baaz:
+                baaz
+            }
+        """
+
+        testFormatting(for: input, output, rule: FormatRules.indent)
+    }
+
+    func testIndentIfExpressionWithComments() {
+        let input = """
+        let foo =
+            // There is a comment before the first branch
+            if let foo {
+                foo
+            }
+            // And also a comment before the second branch
+            else {
+                bar
+            }
+
+        let foo =
+            // There is a multi-line comment before the first branch,
+            // check it out I am a multi-line comment
+            if let foo {
+                foo
+            }
+            //
+            // There is a multi-line comment before the second branch too
+            //
+            else {
+                bar
+            }
+        """
+
+        //
+        // TODO: Fix issue with wrapMultilineStatementBraces
+        //
+        testFormatting(for: input, rule: FormatRules.indent, exclude: ["wrapMultilineStatementBraces"])
+    }
+
+    func testSE0380Example() {
+        let input = """
+        let bullet =
+            if isRoot && (count == 0 || !willExpand) { "" }
+            else if count == 0 { "- " }
+            else if maxDepth <= 0 { "▹ " }
+            else { "▿ " }
+
+        print(bullet)
+        """
+        let options = FormatOptions()
+        testFormatting(for: input, rule: FormatRules.indent, options: options, exclude: ["wrapConditionalBodies", "andOperator", "redundantParens"])
+    }
 }

Tests/RulesTests+Redundancy.swift

@@ -1273,7 +1273,7 @@ class RedundancyTests: RulesTests {
         }
         """
         let options = FormatOptions(redundantType: .inferred, swiftVersion: "5.9")
-        testFormatting(for: input, rule: FormatRules.redundantType, options: options)
+        testFormatting(for: input, rule: FormatRules.redundantType, options: options, exclude: ["wrapMultilineConditionalAssignment"])
     }
 
     func testRedundantTypeWithIfExpression_inferred() {
@@ -1292,7 +1292,7 @@ class RedundancyTests: RulesTests {
         }
         """
         let options = FormatOptions(redundantType: .inferred, swiftVersion: "5.9")
-        testFormatting(for: input, output, rule: FormatRules.redundantType, options: options)
+        testFormatting(for: input, output, rule: FormatRules.redundantType, options: options, exclude: ["wrapMultilineConditionalAssignment"])
     }
 
     func testRedundantTypeWithIfExpression_explicit() {
@@ -1311,7 +1311,7 @@ class RedundancyTests: RulesTests {
         }
         """
         let options = FormatOptions(redundantType: .explicit, swiftVersion: "5.9")
-        testFormatting(for: input, output, rule: FormatRules.redundantType, options: options)
+        testFormatting(for: input, output, rule: FormatRules.redundantType, options: options, exclude: ["wrapMultilineConditionalAssignment"])
     }
 
     func testRedundantTypeWithNestedIfExpression_inferred() {
@@ -1348,7 +1348,7 @@ class RedundancyTests: RulesTests {
         }
         """
         let options = FormatOptions(redundantType: .inferred, swiftVersion: "5.9")
-        testFormatting(for: input, output, rule: FormatRules.redundantType, options: options)
+        testFormatting(for: input, output, rule: FormatRules.redundantType, options: options, exclude: ["wrapMultilineConditionalAssignment"])
     }
 
     func testRedundantTypeWithNestedIfExpression_explicit() {
@@ -1385,7 +1385,7 @@ class RedundancyTests: RulesTests {
         }
         """
         let options = FormatOptions(redundantType: .explicit, swiftVersion: "5.9")
-        testFormatting(for: input, output, rule: FormatRules.redundantType, options: options)
+        testFormatting(for: input, output, rule: FormatRules.redundantType, options: options, exclude: ["wrapMultilineConditionalAssignment"])
     }
 
     func testRedundantTypeWithLiteralsInIfExpression() {
@@ -1404,7 +1404,7 @@ class RedundancyTests: RulesTests {
         }
         """
         let options = FormatOptions(redundantType: .inferred, swiftVersion: "5.9")
-        testFormatting(for: input, output, rule: FormatRules.redundantType, options: options)
+        testFormatting(for: input, output, rule: FormatRules.redundantType, options: options, exclude: ["wrapMultilineConditionalAssignment"])
     }
 
     // --redundanttype explicit
@@ -2694,7 +2694,7 @@ class RedundancyTests: RulesTests {
         }
         """
         let options = FormatOptions(swiftVersion: "5.9")
-        testFormatting(for: input, [output], rules: [FormatRules.redundantReturn, FormatRules.redundantClosure, FormatRules.indent], options: options)
+        testFormatting(for: input, [output], rules: [FormatRules.redundantReturn, FormatRules.redundantClosure, FormatRules.indent], options: options, exclude: ["wrapMultilineConditionalAssignment"])
     }
 
     func testRedundantSwitchStatementReturnInFunction() {
@@ -7850,7 +7850,7 @@ class RedundancyTests: RulesTests {
 
         let options = FormatOptions(swiftVersion: "5.9")
         testFormatting(for: input, [output], rules: [FormatRules.redundantReturn, FormatRules.redundantClosure],
-                       options: options, exclude: ["indent"])
+                       options: options, exclude: ["indent", "wrapMultilineConditionalAssignment"])
     }
 
     func testRedundantClosureWithExplicitReturn2() {
@@ -8247,7 +8247,7 @@ class RedundancyTests: RulesTests {
         """
         let options = FormatOptions(swiftVersion: "5.9")
         testFormatting(for: input, [output], rules: [FormatRules.redundantReturn, FormatRules.redundantClosure],
-                       options: options, exclude: ["indent"])
+                       options: options, exclude: ["indent", "wrapMultilineConditionalAssignment"])
     }
 
     func testRedundantClosureDoesntLeaveStrayTryAwait() {
@@ -8269,7 +8269,7 @@ class RedundancyTests: RulesTests {
         """
         let options = FormatOptions(swiftVersion: "5.9")
         testFormatting(for: input, [output], rules: [FormatRules.redundantReturn, FormatRules.redundantClosure],
-                       options: options, exclude: ["indent"])
+                       options: options, exclude: ["indent", "wrapMultilineConditionalAssignment"])
     }
 
     func testRedundantClosureDoesntLeaveInvalidSwitchExpressionInOperatorChain() {
@@ -8387,7 +8387,7 @@ class RedundancyTests: RulesTests {
         """
 
         let options = FormatOptions(swiftVersion: "5.9")
-        testFormatting(for: input, output, rule: FormatRules.redundantClosure, options: options, exclude: ["indent"])
+        testFormatting(for: input, output, rule: FormatRules.redundantClosure, options: options, exclude: ["indent", "wrapMultilineConditionalAssignment"])
     }
 
     func testRedundantClosureDoesntLeaveInvalidSwitchExpressionInArray() {
@@ -8722,7 +8722,7 @@ class RedundancyTests: RulesTests {
         """
 
         let options = FormatOptions(swiftVersion: "5.10")
-        testFormatting(for: input, output, rule: FormatRules.redundantClosure, options: options, exclude: ["indent"])
+        testFormatting(for: input, output, rule: FormatRules.redundantClosure, options: options, exclude: ["indent", "wrapMultilineConditionalAssignment"])
     }
 
     func testRedundantClosureDoesntBreakBuildWithRedundantReturnRuleDisabled() {
@@ -8785,7 +8785,7 @@ class RedundancyTests: RulesTests {
 
         let options = FormatOptions(swiftVersion: "5.9")
         testFormatting(for: input, [output], rules: [FormatRules.redundantReturn, FormatRules.redundantClosure],
-                       options: options, exclude: ["indent", "blankLinesBetweenScopes"])
+                       options: options, exclude: ["indent", "blankLinesBetweenScopes", "wrapMultilineConditionalAssignment"])
     }
 
     // MARK: - redundantOptionalBinding