Fix formatting of keypath subscript components with labels.

Also refine the treatment of the closing bracket so that it is only
forced to go onto a new line if the subscript component is not the
last component in the keypath. This aligns with how we handle
closing brackets/parentheses in member access chains.

Fixes #663.

Author: allevato

Sources/SwiftFormat/PrettyPrint/TokenStreamCreator.swift

@@ -1894,11 +1894,31 @@ fileprivate final class TokenStreamCreator: SyntaxVisitor {
   }
 
   override func visit(_ node: KeyPathSubscriptComponentSyntax) -> SyntaxVisitorContinueKind {
-    after(node.leftSquare, tokens: .break(.open, size: 0), .open)
-    before(node.rightSquare, tokens: .break(.close, size: 0), .close)
+    var breakBeforeRightParen = !isCompactSingleFunctionCallArgument(node.arguments)
+    if let component = node.parent?.as(KeyPathComponentSyntax.self) {
+      breakBeforeRightParen = !isLastKeyPathComponent(component)
+    }
+
+    arrangeFunctionCallArgumentList(
+      node.arguments,
+      leftDelimiter: node.leftSquare,
+      rightDelimiter: node.rightSquare,
+      forcesBreakBeforeRightDelimiter: breakBeforeRightParen)
     return .visitChildren
   }
 
+  /// Returns a value indicating whether the given key path component was the last component in the
+  /// list containing it.
+  private func isLastKeyPathComponent(_ component: KeyPathComponentSyntax) -> Bool {
+    guard
+      let componentList = component.parent?.as(KeyPathComponentListSyntax.self),
+      let lastComponent = componentList.last
+    else {
+      return false
+    }
+    return component == lastComponent
+  }
+
   override func visit(_ node: TernaryExprSyntax) -> SyntaxVisitorContinueKind {
     // 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`

Tests/SwiftFormatTests/PrettyPrint/KeyPathExprTests.swift

@@ -131,7 +131,9 @@ final class KeyPathExprTests: PrettyPrintTestCase {
       #"""
       let x = \ReallyLongType.reallyLongProperty.anotherLongProperty
       let x = \.reeeeallyLongProperty.anotherLongProperty
+      let x = \.longProperty.a.b.c[really + long + expression]
       let x = \.longProperty.a.b.c[really + long + expression].anotherLongProperty
+      let x = \.longProperty.a.b.c[label:really + long + expression].anotherLongProperty
       """#
 
     let expected =
@@ -146,6 +148,16 @@ final class KeyPathExprTests: PrettyPrintTestCase {
       let x =
         \.longProperty.a.b.c[
           really + long
+            + expression]
+      let x =
+        \.longProperty.a.b.c[
+          really + long
+            + expression
+        ].anotherLongProperty
+      let x =
+        \.longProperty.a.b.c[
+          label: really
+            + long
             + expression
         ].anotherLongProperty