Merge pull request swiftlang#108 from allevato/nested-ternaries

Vastly improve formatting of ternary expressions.

Author: allevato

Sources/SwiftFormatPrettyPrint/TokenStreamCreator.swift

@@ -1256,11 +1256,19 @@ private final class TokenStreamCreator: SyntaxVisitor {
   }
 
   func visit(_ node: TernaryExprSyntax) -> SyntaxVisitorContinueKind {
-    before(node.questionMark, tokens: .break, .open)
+    // The order of the .open/.close tokens here is intentional. They are normally paired with the
+    // corresponding breaks, but in this case, we want to prioritize keeping the entire `? a : b`
+    // part together if some part of the ternary wraps, instead of keeping `c ? a` together and
+    // wrapping after that.
+    before(node.questionMark, tokens: .break(.open(kind: .continuation)), .open)
     after(node.questionMark, tokens: .space)
-    before(node.colonMark, tokens: .break, .open)
+    before(
+      node.colonMark,
+      tokens: .break(.close(mustBreak: false), size: 0), .break(.open(kind: .continuation)), .open)
     after(node.colonMark, tokens: .space)
-    after(node.secondChoice.lastToken, tokens: .close, .close)
+    after(
+      node.secondChoice.lastToken,
+      tokens: .break(.close(mustBreak: false), size: 0), .close, .close)
     return .visitChildren
   }
 
@@ -2540,6 +2548,8 @@ private final class TokenStreamCreator: SyntaxVisitor {
     switch expr {
     case let sequenceExpr as SequenceExprSyntax:
       return sequenceExpr.elements.count > 1
+    case is TernaryExprSyntax:
+      return true
     case let tryExpr as TryExprSyntax:
       return isCompoundExpression(tryExpr.expression)
     case let tupleExpr as TupleExprSyntax where tupleExpr.elementList.count == 1:
@@ -2619,6 +2629,12 @@ private final class TokenStreamCreator: SyntaxVisitor {
       }
     }
 
+    // If the right-hand-side is a ternary expression, stack indentation around the condition so
+    // that it is indented relative to the `?` and `:` tokens.
+    if let ternaryExpr = rhs as? TernaryExprSyntax {
+      return (unindentingNode: ternaryExpr.conditionExpression, shouldReset: false)
+    }
+
     // If the right-hand-side of the operator is or starts with a parenthesized expression, stack
     // indentation around the operator and those parentheses. We don't need to reset here because
     // the parentheses are sufficient to provide a visual indication of the nesting relationship.

Tests/SwiftFormatPrettyPrintTests/TernaryExprTests.swift

@@ -16,14 +16,17 @@ public class TernaryExprTests: PrettyPrintTestCase {
       let x = a ? b : c
       let y = a ? b : c
       let z = a ? b : c
-      let reallyLongName = a
-        ? longTruePart : longFalsePart
-      let reallyLongName = reallyLongCondition
+      let reallyLongName =
+        a ? longTruePart : longFalsePart
+      let reallyLongName =
+        reallyLongCondition
         ? reallyLongTruePart : reallyLongFalsePart
-      let reallyLongName = reallyLongCondition
+      let reallyLongName =
+        reallyLongCondition
         ? reallyReallyReallyLongTruePart
         : reallyLongFalsePart
-      let reallyLongName = someCondition
+      let reallyLongName =
+        someCondition
         ? firstFunc(foo: arg)
         : secondFunc(bar: arg)
 
@@ -52,4 +55,156 @@ public class TernaryExprTests: PrettyPrintTestCase {
 
     assertPrettyPrintEqual(input: input, expected: expected, linelength: 80)
   }
+
+  public func testTernaryWithWrappingExpressions() {
+    let input =
+      """
+      foo = firstTerm + secondTerm + thirdTerm ? firstTerm + secondTerm + thirdTerm : firstTerm + secondTerm + thirdTerm
+      let foo = firstTerm + secondTerm + thirdTerm ? firstTerm + secondTerm + thirdTerm : firstTerm + secondTerm + thirdTerm
+      foo = firstTerm || secondTerm && thirdTerm ? firstTerm + secondTerm + thirdTerm : firstTerm + secondTerm + thirdTerm
+      let foo = firstTerm || secondTerm && thirdTerm ? firstTerm + secondTerm + thirdTerm : firstTerm + secondTerm + thirdTerm
+      """
+
+    let expected =
+      """
+      foo =
+        firstTerm
+          + secondTerm
+          + thirdTerm
+        ? firstTerm
+          + secondTerm
+          + thirdTerm
+        : firstTerm
+          + secondTerm
+          + thirdTerm
+      let foo =
+        firstTerm
+          + secondTerm
+          + thirdTerm
+        ? firstTerm
+          + secondTerm
+          + thirdTerm
+        : firstTerm
+          + secondTerm
+          + thirdTerm
+      foo =
+        firstTerm
+          || secondTerm
+            && thirdTerm
+        ? firstTerm
+          + secondTerm
+          + thirdTerm
+        : firstTerm
+          + secondTerm
+          + thirdTerm
+      let foo =
+        firstTerm
+          || secondTerm
+            && thirdTerm
+        ? firstTerm
+          + secondTerm
+          + thirdTerm
+        : firstTerm
+          + secondTerm
+          + thirdTerm
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 15)
+  }
+
+  public func testNestedTernaries() {
+    let input =
+      """
+      a = b ? c : d ? e : f
+      let a = b ? c : d ? e : f
+      a = b ? c0 + c1 : d ? e : f
+      let a = b ? c0 + c1 : d ? e : f
+      a = b ? c0 + c1 + c2 + c3 : d ? e : f
+      let a = b ? c0 + c1 + c2 + c3 : d ? e : f
+      foo = testA ? testB ? testC : testD : testE ? testF : testG
+      let foo = testA ? testB ? testC : testD : testE ? testF : testG
+      """
+
+    let expected =
+      """
+      a =
+        b
+        ? c
+        : d ? e : f
+      let a =
+        b
+        ? c
+        : d ? e : f
+      a =
+        b
+        ? c0 + c1
+        : d ? e : f
+      let a =
+        b
+        ? c0 + c1
+        : d ? e : f
+      a =
+        b
+        ? c0 + c1
+          + c2 + c3
+        : d ? e : f
+      let a =
+        b
+        ? c0 + c1
+          + c2 + c3
+        : d ? e : f
+      foo =
+        testA
+        ? testB
+          ? testC
+          : testD
+        : testE
+          ? testF
+          : testG
+      let foo =
+        testA
+        ? testB
+          ? testC
+          : testD
+        : testE
+          ? testF
+          : testG
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 15)
+  }
+
+  public func testExpressionStartsWithTernary() {
+    // When the ternary itself doesn't already start on a continuation line, we don't have a way
+    // to indent the continuation of the condition differently from the first and second choices,
+    // because we don't want to double-indent the condition's continuation lines, and we don't want
+    // to keep put the choices at the same indentation level as the condition (because that would
+    // be the start of the statement). Neither of these choices is really ideal, unfortunately.
+    let input =
+      """
+      condition ? callSomething() : callSomethingElse()
+      condition && otherCondition ? callSomething() : callSomethingElse()
+      (condition && otherCondition) ? callSomething() : callSomethingElse()
+      """
+
+    let expected =
+      """
+      condition
+        ? callSomething()
+        : callSomethingElse()
+      condition
+        && otherCondition
+        ? callSomething()
+        : callSomethingElse()
+      (condition
+        && otherCondition)
+        ? callSomething()
+        : callSomethingElse()
+
+      """
+
+    assertPrettyPrintEqual(input: input, expected: expected, linelength: 25)
+  }
 }

Tests/SwiftFormatPrettyPrintTests/XCTestManifests.swift

@@ -632,8 +632,11 @@ extension TernaryExprTests {
     //   `swift test --generate-linuxmain`
     // to regenerate.
     static let __allTests__TernaryExprTests = [
+        ("testExpressionStartsWithTernary", testExpressionStartsWithTernary),
+        ("testNestedTernaries", testNestedTernaries),
         ("testTernaryExprs", testTernaryExprs),
         ("testTernaryExprsWithMultiplePartChoices", testTernaryExprsWithMultiplePartChoices),
+        ("testTernaryWithWrappingExpressions", testTernaryWithWrappingExpressions),
     ]
 }