Refine expansion of function-typed placeholders

This addresses <swiftlang/sourcekit-lsp#1788>.

Per <swiftlang/sourcekit-lsp#1789>, these
placeholders no longer expand to multi-line trailing closures. Instead
they become inline closures, with nested placeholders for the signature
and body.

This introduces a new formatter, `ClosureLiteralFormat`, so that other
uses of `BasicFormat` are not impacted. `ClosureLiteralFormat` does not
insert newlines into a closure when there is only a single statement in
the body.

Author: woolsweater

Sources/SwiftBasicFormat/ClosureLiteralFormat.swift

@@ -0,0 +1,65 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#if swift(>=6)
+@_spi(RawSyntax) public import SwiftSyntax
+#else
+@_spi(RawSyntax) import SwiftSyntax
+#endif
+
+/// A specialization of `BasicFormat` for closure literals, which is more
+/// conservative about newline insertion.
+///
+/// A closure that seems to be a simple predicate or transform — based on the
+/// number of enclosed statements — will not be reformatted to multiple lines.
+open class ClosureLiteralFormat: BasicFormat {
+  open override func requiresNewline(between first: TokenSyntax?, and second: TokenSyntax?) -> Bool {
+    if let first, isEndOfSmallClosureSignature(first) {
+      return false
+    } else if let first, isSmallClosureDelimiter(first, kind: \.leftBrace) {
+      return false
+    } else if let second, isSmallClosureDelimiter(second, kind: \.rightBrace) {
+      return false
+    } else {
+      return super.requiresNewline(between: first, and: second)
+    }
+  }
+
+  /// Returns `true` if `token` is an opening or closing brace (according to
+  /// `kind`) of a closure, and that closure has no more than one statement in
+  /// its body.
+  private func isSmallClosureDelimiter(
+    _ token: TokenSyntax,
+    kind: KeyPath<ClosureExprSyntax, TokenSyntax>
+  ) -> Bool {
+    guard token.keyPathInParent == kind,
+      let closure = token.parent?.as(ClosureExprSyntax.self)
+    else {
+      return false
+    }
+
+    return closure.statements.count <= 1
+  }
+
+  /// Returns `true` if `token` is the last token in the signature of a closure,
+  /// and that closure has no more than one statement in its body.
+  private func isEndOfSmallClosureSignature(_ token: TokenSyntax) -> Bool {
+    guard let signature = token.ancestorOrSelf(mapping: { $0.as(ClosureSignatureSyntax.self) }),
+      let closure = signature.parent?.as(ClosureExprSyntax.self)
+    else {
+      return false
+    }
+
+    return signature.lastToken(viewMode: viewMode) == token
+      && closure.statements.count <= 1
+  }
+}

Sources/SwiftRefactor/ExpandEditorPlaceholder.swift

@@ -50,9 +50,7 @@ import SwiftSyntaxBuilder
 ///
 /// ### After
 /// ```swift
-/// { someInt in
-///   <#T##String#>
-/// }
+/// <#{ <#someInt#> in <#T##String#> }#>
 /// ```
 ///
 /// ## Other Type Placeholder
@@ -94,18 +92,18 @@ struct ExpandSingleEditorPlaceholder: EditRefactoringProvider {
 
     let expanded: String
     if let functionType = placeholder.typeForExpansion?.as(FunctionTypeSyntax.self) {
-      let basicFormat = BasicFormat(
+      let format = ClosureLiteralFormat(
         indentationWidth: context.indentationWidth,
         initialIndentation: context.initialIndentation
       )
-      var formattedExpansion = functionType.closureExpansion.formatted(using: basicFormat).description
+      var formattedExpansion = functionType.closureExpansion.formatted(using: format).description
       // Strip the initial indentation from the placeholder itself. We only introduced the initial indentation to
       // format consecutive lines. We don't want it at the front of the initial line because it replaces an expression
       // that might be in the middle of a line.
       if formattedExpansion.hasPrefix(context.initialIndentation.description) {
         formattedExpansion = String(formattedExpansion.dropFirst(context.initialIndentation.description.count))
       }
-      expanded = formattedExpansion
+      expanded = wrapInPlaceholder(formattedExpansion)
     } else {
       expanded = placeholder.displayText
     }
@@ -117,9 +115,9 @@ struct ExpandSingleEditorPlaceholder: EditRefactoringProvider {
 }
 
 /// If a function-typed placeholder is the argument to a non-trailing closure
-/// call, expands it and any adjacent function-typed placeholders to trailing
-/// closures on that call. All other placeholders will expand as per
-/// `ExpandEditorPlaceholder`.
+/// call, expands it and any adjacent function-typed placeholders to literal
+/// closures with inner placeholders on that call. All other placeholders will
+/// expand as per `ExpandEditorPlaceholder`.
 ///
 /// ## Before
 /// ```swift
@@ -137,12 +135,10 @@ struct ExpandSingleEditorPlaceholder: EditRefactoringProvider {
 /// foo(
 ///   closure1: <#T##(Int) -> String##(Int) -> String##(_ someInt: Int) -> String#>,
 ///   normalArg: <#T##Int#>,
-///   closure2: { ... }
-/// ) { someInt in
-///   <#T##String#>
-/// } closure2: { someInt in
-///   <#T##String#>
-/// }
+///   closure2: { ... },
+///   closure3: { <#someInt#> in <#T##String#> },
+///   closure4: { <#someInt#> in <#T##String#> }
+/// )
 /// ```
 ///
 /// Expansion on `closure1` and `normalArg` is the same as `ExpandSingleEditorPlaceholder`.
@@ -161,7 +157,7 @@ public struct ExpandEditorPlaceholder: EditRefactoringProvider {
       let arg = placeholder.parent?.as(LabeledExprSyntax.self),
       let argList = arg.parent?.as(LabeledExprListSyntax.self),
       let call = argList.parent?.as(FunctionCallExprSyntax.self),
-      let expandedTrailingClosures = ExpandEditorPlaceholdersToTrailingClosures.expandTrailingClosurePlaceholders(
+      let expandedClosures = ExpandEditorPlaceholdersToLiteralClosures.expandClosurePlaceholders(
         in: call,
         ifIncluded: arg,
         indentationWidth: context.indentationWidth
@@ -170,11 +166,11 @@ public struct ExpandEditorPlaceholder: EditRefactoringProvider {
       return ExpandSingleEditorPlaceholder.textRefactor(syntax: token)
     }
 
-    return [SourceEdit.replace(call, with: expandedTrailingClosures.description)]
+    return [SourceEdit.replace(call, with: expandedClosures.description)]
   }
 }
 
-/// Expand all the editor placeholders in the function call that can be converted to trailing closures.
+/// Expand all the editor placeholders in the function call to literal closures.
 ///
 /// ## Before
 /// ```swift
@@ -189,13 +185,11 @@ public struct ExpandEditorPlaceholder: EditRefactoringProvider {
 /// ```swift
 /// foo(
 ///   arg: <#T##Int#>,
-/// ) { someInt in
-///   <#T##String#>
-/// } secondClosure: { someInt in
-///   <#T##String#>
-/// }
+///   firstClosure: { <#someInt#> in <#T##String#> },
+///   secondClosure: { <#someInt#> in <#T##String#> }
+/// )
 /// ```
-public struct ExpandEditorPlaceholdersToTrailingClosures: SyntaxRefactoringProvider {
+public struct ExpandEditorPlaceholdersToLiteralClosures: SyntaxRefactoringProvider {
   public struct Context {
     public let indentationWidth: Trivia?
 
@@ -208,32 +202,24 @@ public struct ExpandEditorPlaceholdersToTrailingClosures: SyntaxRefactoringProvi
     syntax call: FunctionCallExprSyntax,
     in context: Context = Context()
   ) -> FunctionCallExprSyntax? {
-    return Self.expandTrailingClosurePlaceholders(in: call, ifIncluded: nil, indentationWidth: context.indentationWidth)
+    return Self.expandClosurePlaceholders(
+      in: call,
+      ifIncluded: nil,
+      indentationWidth: context.indentationWidth
+    )
   }
 
   /// If the given argument is `nil` or one of the last arguments that are all
   /// function-typed placeholders and this call doesn't have a trailing
   /// closure, then return a replacement of this call with one that uses
   /// closures based on the function types provided by each editor placeholder.
   /// Otherwise return nil.
-  fileprivate static func expandTrailingClosurePlaceholders(
+  fileprivate static func expandClosurePlaceholders(
     in call: FunctionCallExprSyntax,
     ifIncluded arg: LabeledExprSyntax?,
     indentationWidth: Trivia?
   ) -> FunctionCallExprSyntax? {
-    guard let expanded = call.expandTrailingClosurePlaceholders(ifIncluded: arg, indentationWidth: indentationWidth)
-    else {
-      return nil
-    }
-
-    let callToTrailingContext = CallToTrailingClosures.Context(
-      startAtArgument: call.arguments.count - expanded.numClosures
-    )
-    guard let trailing = CallToTrailingClosures.refactor(syntax: expanded.expr, in: callToTrailingContext) else {
-      return nil
-    }
-
-    return trailing
+    return call.expandClosurePlaceholders(ifIncluded: arg, indentationWidth: indentationWidth)
   }
 }
 
@@ -244,9 +230,7 @@ extension FunctionTypeSyntax {
   /// ```
   /// would become
   /// ```
-  /// { someInt in
-  ///   <#T##String#>
-  /// }
+  /// { <#someInt#> in <#T##String#> }
   /// ```
   fileprivate var closureExpansion: ClosureExprSyntax {
     let closureSignature: ClosureSignatureSyntax?
@@ -311,10 +295,10 @@ extension FunctionCallExprSyntax {
   /// closure, then return a replacement of this call with one that uses
   /// closures based on the function types provided by each editor placeholder.
   /// Otherwise return nil.
-  fileprivate func expandTrailingClosurePlaceholders(
+  fileprivate func expandClosurePlaceholders(
     ifIncluded: LabeledExprSyntax?,
     indentationWidth: Trivia?
-  ) -> (expr: FunctionCallExprSyntax, numClosures: Int)? {
+  ) -> FunctionCallExprSyntax? {
     var includedArg = false
     var argsToExpand = 0
     for arg in arguments.reversed() {
@@ -359,10 +343,7 @@ extension FunctionCallExprSyntax {
     }
 
     let originalArgs = arguments.dropLast(argsToExpand)
-    return (
-      detached.with(\.arguments, LabeledExprListSyntax(originalArgs + expandedArgs)),
-      expandedArgs.count
-    )
+    return detached.with(\.arguments, LabeledExprListSyntax(originalArgs + expandedArgs))
   }
 }
 

Sources/SwiftSyntax/SyntaxProtocol.swift

@@ -235,7 +235,10 @@ extension SyntaxProtocol {
     return self.previousToken(viewMode: .sourceAccurate)
   }
 
-  /// Returns this node or the first ancestor that satisfies `condition`.
+  /// Applies `map` to this node and each of its ancestors until a non-`nil`
+  /// value is produced, then returns that value.
+  ///
+  /// If no node has a non-`nil` mapping, returns `nil`.
   public func ancestorOrSelf<T>(mapping map: (Syntax) -> T?) -> T? {
     return self.withUnownedSyntax {
       var node = $0

Tests/SwiftBasicFormatTest/ClosureLiteralFormatTests.swift

@@ -0,0 +1,129 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SwiftBasicFormat
+import SwiftParser
+import SwiftSyntax
+@_spi(Testing) import SwiftSyntaxBuilder
+import XCTest
+import _SwiftSyntaxTestSupport
+
+fileprivate func assertFormatted<T: SyntaxProtocol>(
+  tree: T,
+  expected: String,
+  using format: ClosureLiteralFormat = ClosureLiteralFormat(indentationWidth: .spaces(4)),
+  file: StaticString = #filePath,
+  line: UInt = #line
+) {
+  assertStringsEqualWithDiff(tree.formatted(using: format).description, expected, file: file, line: line)
+}
+
+fileprivate func assertFormatted(
+  source: String,
+  expected: String,
+  using format: ClosureLiteralFormat = ClosureLiteralFormat(indentationWidth: .spaces(4)),
+  file: StaticString = #filePath,
+  line: UInt = #line
+) {
+  assertFormatted(
+    tree: Parser.parse(source: source),
+    expected: expected,
+    using: format,
+    file: file,
+    line: line
+  )
+}
+
+final class ClosureLiteralFormatTests: XCTestCase {
+  func testSingleStatementClosureArg() {
+    assertFormatted(
+      source: """
+        foo(bar: { baz in baz.quux })
+        """,
+      expected: """
+        foo(bar: { baz in baz.quux })
+        """
+    )
+  }
+
+  func testSingleStatementClosureArgAlreadyMultiLine() {
+    assertFormatted(
+      source: """
+        foo(
+          bar: { baz in
+            baz.quux
+          }
+        )
+        """,
+      expected: """
+        foo(
+          bar: { baz in
+            baz.quux
+          }
+        )
+        """
+    )
+  }
+
+  func testMultiStatmentClosureArg() {
+    assertFormatted(
+      source: """
+        foo(
+            bar: { baz in print(baz); return baz.quux }
+        )
+        """,
+      expected: """
+        foo(
+            bar: { baz in
+                print(baz);
+                return baz.quux
+            }
+        )
+        """
+    )
+  }
+
+  func testMultiStatementClosureArgAlreadyMultiLine() {
+    assertFormatted(
+      source: """
+        foo(
+            bar: { baz in
+                print(baz)
+                return baz.quux
+            }
+        )
+        """,
+      expected: """
+        foo(
+            bar: { baz in
+                print(baz)
+                return baz.quux
+            }
+        )
+        """
+    )
+  }
+
+  func testFormatClosureWithInitialIndentation() throws {
+    assertFormatted(
+      tree: ClosureExprSyntax(
+        statements: CodeBlockItemListSyntax([
+          CodeBlockItemSyntax(item: CodeBlockItemSyntax.Item(IntegerLiteralExprSyntax(integerLiteral: 2)))
+        ])
+      ),
+      expected: """
+            { 2 }
+        """,
+      using: ClosureLiteralFormat(initialIndentation: .spaces(4))
+    )
+  }
+}