Merge pull request #553 from allevato/if-switch-exprs

Improve wrapping of if/switch expressions.

Author: allevato

Sources/SwiftFormatPrettyPrint/TokenStreamCreator.swift

@@ -1965,13 +1965,18 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
 
         // If the rhs starts with a parenthesized expression, stack indentation around it.
         // Otherwise, use regular continuation breaks.
-        if let (unindentingNode, _, breakKind, _) =
+        if let (unindentingNode, _, breakKind, shouldGroup) =
           stackedIndentationBehavior(after: binOp, rhs: rhs)
         {
           beforeTokens = [.break(.open(kind: breakKind))]
+          var afterTokens: [Token] = [.break(.close(mustBreak: false), size: 0)]
+          if shouldGroup {
+            beforeTokens.append(.open)
+            afterTokens.append(.close)
+          }
           after(
             unindentingNode.lastToken(viewMode: .sourceAccurate),
-            tokens: [.break(.close(mustBreak: false), size: 0)])
+            tokens: afterTokens)
         } else {
           beforeTokens = [.break(.continue)]
         }
@@ -2121,9 +2126,17 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
     if let initializer = node.initializer {
       let expr = initializer.value
 
-      if let (unindentingNode, _, breakKind, _) = stackedIndentationBehavior(rhs: expr) {
-        after(initializer.equal, tokens: .break(.open(kind: breakKind)))
-        after(unindentingNode.lastToken(viewMode: .sourceAccurate), tokens: .break(.close(mustBreak: false), size: 0))
+      if let (unindentingNode, _, breakKind, shouldGroup) = stackedIndentationBehavior(rhs: expr) {
+        var openTokens: [Token] = [.break(.open(kind: breakKind))]
+        if shouldGroup {
+          openTokens.append(.open)
+        }
+        after(initializer.equal, tokens: openTokens)
+        var closeTokens: [Token] = [.break(.close(mustBreak: false), size: 0)]
+        if shouldGroup {
+          closeTokens.append(.close)
+        }
+        after(unindentingNode.lastToken(viewMode: .sourceAccurate), tokens: closeTokens)
       } else {
         after(initializer.equal, tokens: .break(.continue))
       }
@@ -2290,9 +2303,13 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
   override func visit(_ node: InitializerClauseSyntax) -> SyntaxVisitorContinueKind {
     before(node.equal, tokens: .space)
 
-    // InitializerClauses that are children of a PatternBindingSyntax are already handled in the
-    // latter node, to ensure that continuations stack appropriately.
-    if node.parent == nil || !node.parent!.is(PatternBindingSyntax.self) {
+    // InitializerClauses that are children of a PatternBindingSyntax or
+    // OptionalBindingConditionSyntax are already handled in the latter node, to ensure that
+    // continuations stack appropriately.
+    if let parent = node.parent,
+      !parent.is(PatternBindingSyntax.self)
+        && !parent.is(OptionalBindingConditionSyntax.self)
+    {
       after(node.equal, tokens: .break)
     }
     return .visitChildren
@@ -2474,6 +2491,26 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
       after(typeAnnotation.lastToken(viewMode: .sourceAccurate), tokens: .break(.close(mustBreak: false), size: 0))
     }
 
+    if let initializer = node.initializer {
+      if let (unindentingNode, _, breakKind, shouldGroup) =
+        stackedIndentationBehavior(rhs: initializer.value)
+      {
+        var openTokens: [Token] = [.break(.open(kind: breakKind))]
+        if shouldGroup {
+          openTokens.append(.open)
+        }
+        after(initializer.equal, tokens: openTokens)
+
+        var closeTokens: [Token] = [.break(.close(mustBreak: false), size: 0)]
+        if shouldGroup {
+          closeTokens.append(.close)
+        }
+        after(unindentingNode.lastToken(viewMode: .sourceAccurate), tokens: closeTokens)
+      } else {
+        after(initializer.equal, tokens: .break(.continue))
+      }
+    }
+
     return .visitChildren
   }
 
@@ -3389,47 +3426,63 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
     }
   }
 
-  /// Walks the expression and returns the leftmost multiline string literal (which might be the
-  /// expression itself) if the leftmost child is a multiline string literal or if it is a unary
-  /// operation applied to a multiline string literal.
+  /// Walks the expression and returns the leftmost subexpression (which might be the expression
+  /// itself) if the leftmost child is a node of the given type or if it is a unary operation
+  /// applied to a node of the given type.
   ///
-  /// - Parameter expr: The expression whose leftmost multiline string literal should be returned.
-  /// - Returns: The leftmost multiline string literal, or nil if the leftmost subexpression was
-  ///   not a multiline string literal.
-  private func leftmostMultilineStringLiteral(of expr: ExprSyntax) -> StringLiteralExprSyntax? {
+  /// - Parameter expr: The expression whose leftmost matching subexpression should be returned.
+  /// - Returns: The leftmost subexpression, or nil if the leftmost subexpression was not the
+  ///   desired type.
+  private func leftmostExpr(
+    of expr: ExprSyntax,
+    ifMatching predicate: (ExprSyntax) -> Bool
+  ) -> ExprSyntax? {
+    if predicate(expr) {
+      return expr
+    }
     switch Syntax(expr).as(SyntaxEnum.self) {
-    case .stringLiteralExpr(let stringLiteralExpr)
-      where stringLiteralExpr.openQuote.tokenKind == .multilineStringQuote:
-      return stringLiteralExpr
     case .infixOperatorExpr(let infixOperatorExpr):
-      return leftmostMultilineStringLiteral(of: infixOperatorExpr.leftOperand)
+      return leftmostExpr(of: infixOperatorExpr.leftOperand, ifMatching: predicate)
     case .asExpr(let asExpr):
-      return leftmostMultilineStringLiteral(of: asExpr.expression)
+      return leftmostExpr(of: asExpr.expression, ifMatching: predicate)
     case .isExpr(let isExpr):
-      return leftmostMultilineStringLiteral(of: isExpr.expression)
+      return leftmostExpr(of: isExpr.expression, ifMatching: predicate)
     case .forcedValueExpr(let forcedValueExpr):
-      return leftmostMultilineStringLiteral(of: forcedValueExpr.expression)
+      return leftmostExpr(of: forcedValueExpr.expression, ifMatching: predicate)
     case .optionalChainingExpr(let optionalChainingExpr):
-      return leftmostMultilineStringLiteral(of: optionalChainingExpr.expression)
+      return leftmostExpr(of: optionalChainingExpr.expression, ifMatching: predicate)
     case .postfixUnaryExpr(let postfixUnaryExpr):
-      return leftmostMultilineStringLiteral(of: postfixUnaryExpr.expression)
+      return leftmostExpr(of: postfixUnaryExpr.expression, ifMatching: predicate)
     case .prefixOperatorExpr(let prefixOperatorExpr):
-      return leftmostMultilineStringLiteral(of: prefixOperatorExpr.postfixExpression)
+      return leftmostExpr(of: prefixOperatorExpr.postfixExpression, ifMatching: predicate)
     case .ternaryExpr(let ternaryExpr):
-      return leftmostMultilineStringLiteral(of: ternaryExpr.conditionExpression)
+      return leftmostExpr(of: ternaryExpr.conditionExpression, ifMatching: predicate)
     case .functionCallExpr(let functionCallExpr):
-      return leftmostMultilineStringLiteral(of: functionCallExpr.calledExpression)
+      return leftmostExpr(of: functionCallExpr.calledExpression, ifMatching: predicate)
     case .subscriptExpr(let subscriptExpr):
-      return leftmostMultilineStringLiteral(of: subscriptExpr.calledExpression)
+      return leftmostExpr(of: subscriptExpr.calledExpression, ifMatching: predicate)
     case .memberAccessExpr(let memberAccessExpr):
-      return memberAccessExpr.base.flatMap { leftmostMultilineStringLiteral(of: $0) }
+      return memberAccessExpr.base.flatMap { leftmostExpr(of: $0, ifMatching: predicate) }
     case .postfixIfConfigExpr(let postfixIfConfigExpr):
-      return postfixIfConfigExpr.base.flatMap { leftmostMultilineStringLiteral(of: $0) }
+      return postfixIfConfigExpr.base.flatMap { leftmostExpr(of: $0, ifMatching: predicate) }
     default:
       return nil
     }
   }
 
+  /// Walks the expression and returns the leftmost multiline string literal (which might be the
+  /// expression itself) if the leftmost child is a multiline string literal or if it is a unary
+  /// operation applied to a multiline string literal.
+  ///
+  /// - Parameter expr: The expression whose leftmost multiline string literal should be returned.
+  /// - Returns: The leftmost multiline string literal, or nil if the leftmost subexpression was
+  ///   not a multiline string literal.
+  private func leftmostMultilineStringLiteral(of expr: ExprSyntax) -> StringLiteralExprSyntax? {
+    return leftmostExpr(of: expr) {
+      $0.as(StringLiteralExprSyntax.self)?.openQuote.tokenKind == .multilineStringQuote
+    }?.as(StringLiteralExprSyntax.self)
+  }
+
   /// Returns the outermost node enclosing the given node whose closing delimiter(s) must be kept
   /// alongside the last token of the given node. Any tokens between `node.lastToken` and the
   /// returned node's `lastToken` are delimiter tokens that shouldn't be preceded by a break.
@@ -3520,7 +3573,7 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
           unindentingNode: unindentingParenExpr,
           shouldReset: true,
           breakKind: .continuation,
-          shouldGroup: true
+          shouldGroup: false
         )
       }
 
@@ -3536,7 +3589,7 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
         unindentingNode: Syntax(parenthesizedExpr),
         shouldReset: false,
         breakKind: .continuation,
-        shouldGroup: true
+        shouldGroup: false
       )
     }
 
@@ -3552,6 +3605,17 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
       )
     }
 
+    if let leftmostExpr = leftmostExpr(of: rhs, ifMatching: {
+      $0.is(IfExprSyntax.self) || $0.is(SwitchExprSyntax.self)
+    }) {
+      return (
+        unindentingNode: Syntax(leftmostExpr),
+        shouldReset: false,
+        breakKind: .block,
+        shouldGroup: true
+      )
+    }
+
     // Otherwise, don't stack--use regular continuation breaks instead.
     return nil
   }

Tests/SwiftFormatPrettyPrintTests/IfStmtTests.swift

@@ -206,13 +206,14 @@ final class IfStmtTests: PrettyPrintTestCase {
     let expected =
       """
       func foo() -> Int {
-        let x = if var1 < var2 {
-          23
-        } else if d < e {
-          24
-        } else {
-          0
-        }
+        let x =
+          if var1 < var2 {
+            23
+          } else if d < e {
+            24
+          } else {
+            0
+          }
         return x
       }
 
@@ -221,6 +222,39 @@ final class IfStmtTests: PrettyPrintTestCase {
     assertPrettyPrintEqual(input: input, expected: expected, linelength: 26)
   }
 
+  func testIfExpression3() {
+    let input =
+      """
+      let x = if a { b } else { c }
+      xyzab = if a { b } else { c }
+      """
+    assertPrettyPrintEqual(input: input, expected: input + "\n", linelength: 80)
+
+    let expected28 =
+      """
+      let x =
+        if a { b } else { c }
+      xyzab =
+        if a { b } else { c }
+
+      """
+    assertPrettyPrintEqual(input: input, expected: expected28, linelength: 28)
+
+    let expected22 =
+      """
+      let x =
+        if a { b } else {
+          c
+        }
+      xyzab =
+        if a { b } else {
+          c
+        }
+
+      """
+    assertPrettyPrintEqual(input: input, expected: expected22, linelength: 22)
+  }
+
   func testMatchingPatternConditions() {
     let input =
       """

Tests/SwiftFormatPrettyPrintTests/StringTests.swift

@@ -419,6 +419,33 @@ final class StringTests: PrettyPrintTestCase {
     assertPrettyPrintEqual(input: input, expected: input + "\n", linelength: 100)
   }
 
+  func testMultilineStringsNestedInAnotherWrappingContext() {
+    let input =
+      #"""
+      guard
+          let x = """
+              blah
+              blah
+              """.data(using: .utf8) {
+          print(x)
+      }
+      """#
+
+    let expected =
+      #"""
+      guard
+        let x = """
+          blah
+          blah
+          """.data(using: .utf8)
+      {
+        print(x)
+      }
+
+      """#
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 100)
+  }
+
   func testEmptyMultilineStrings() {
     let input =
       ##"""

Tests/SwiftFormatPrettyPrintTests/SwitchStmtTests.swift

@@ -297,20 +297,42 @@ final class SwitchStmtTests: PrettyPrintTestCase {
     let expected =
       """
       func foo() -> Int {
-        let x = switch value1 + value2 + value3 + value4 {
-        case "a":
-          0
-        case "b":
-          1
-        default:
-          2
-        }
+        let x =
+          switch value1 + value2 + value3 + value4 {
+          case "a":
+            0
+          case "b":
+            1
+          default:
+            2
+          }
+        return x
+      }
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 46)
+
+    let expected43 =
+      """
+      func foo() -> Int {
+        let x =
+          switch value1 + value2 + value3
+            + value4
+          {
+          case "a":
+            0
+          case "b":
+            1
+          default:
+            2
+          }
         return x
       }
 
       """
 
-    assertPrettyPrintEqual(input: input, expected: expected, linelength: 52)
+    assertPrettyPrintEqual(input: input, expected: expected43, linelength: 43)
   }
 
   func testUnknownDefault() {