Handle new swift-syntax closure expansion behavior

This resolves <swiftlang#1788>,
following the discussion of alternatives on
<https://github.com/swiftlang/sourcekit-lsp/pulls/1789>. The bulk of the
change updates the translation from SourceKit placeholders to LSP
placeholders to handle nesting.

Sources/SourceKitLSP/Swift/CodeCompletionSession.swift

@@ -323,9 +323,9 @@ class CodeCompletionSession {
     var parser = Parser(exprToExpand)
     let expr = ExprSyntax.parse(from: &parser)
     guard let call = OutermostFunctionCallFinder.findOutermostFunctionCall(in: expr),
-      let expandedCall = ExpandEditorPlaceholdersToTrailingClosures.refactor(
+      let expandedCall = ExpandEditorPlaceholdersToLiteralClosures.refactor(
         syntax: call,
-        in: ExpandEditorPlaceholdersToTrailingClosures.Context(indentationWidth: indentationWidth)
+        in: ExpandEditorPlaceholdersToLiteralClosures.Context(indentationWidth: indentationWidth)
       )
     else {
       return nil
@@ -334,7 +334,7 @@ class CodeCompletionSession {
     let bytesToExpand = Array(exprToExpand.utf8)
 
     var expandedBytes: [UInt8] = []
-    // Add the prefix that we stripped of to allow expression parsing
+    // Add the prefix that we stripped off to allow expression parsing
     expandedBytes += strippedPrefix.utf8
     // Add any part of the expression that didn't end up being part of the function call
     expandedBytes += bytesToExpand[0..<call.position.utf8Offset]

Sources/SourceKitLSP/Swift/RewriteSourceKitPlaceholders.swift

@@ -14,34 +14,163 @@ import Foundation
 import SKLogging
 @_spi(RawSyntax) import SwiftSyntax
 
-func rewriteSourceKitPlaceholders(in string: String, clientSupportsSnippets: Bool) -> String {
-  var result = string
-  var index = 1
-  while let start = result.range(of: "<#") {
-    guard let end = result[start.upperBound...].range(of: "#>") else {
-      logger.fault("Invalid placeholder in \(string)")
-      return string
-    }
-    let rawPlaceholder = String(result[start.lowerBound..<end.upperBound])
-    guard let displayName = nameForSnippet(rawPlaceholder) else {
-      logger.fault("Failed to decode placeholder \(rawPlaceholder) in \(string)")
-      return string
+/// Translate SourceKit placeholder syntax — `<#foo#>` — in `input` to LSP
+/// placeholder syntax: `${n:foo}`.
+///
+/// If `clientSupportsSnippets` is `false`, the placeholder is rendered as an
+/// empty string, to prevent the client from inserting special placeholder
+/// characters as if they were literal text.
+func rewriteSourceKitPlaceholders(in input: String, clientSupportsSnippets: Bool) -> String {
+  var result = ""
+  var nextPlaceholderNumber = 1
+  // Current stack of nested placeholders, most nested last. Each element needs
+  // to be rendered inside the element before it.
+  var placeholders: [(number: Int, contents: String)] = []
+  let tokens = tokenize(input)
+  for token in tokens {
+    switch token {
+    case let .text(text) where placeholders.isEmpty:
+      result += text
+
+    case let .text(text):
+      placeholders.latest.contents += text
+
+    case let .curlyBrace(brace) where placeholders.isEmpty:
+      result.append(brace)
+
+    case let .curlyBrace(brace):
+      // Braces are only escaped _inside_ a placeholder; otherwise the client
+      // would include the backslashes literally.
+      placeholders.latest.contents.append(contentsOf: ["\\", brace])
+
+    case .open:
+      placeholders.append((number: nextPlaceholderNumber, contents: ""))
+      nextPlaceholderNumber += 1
+
+    case .close:
+      guard let (number, placeholderBody) = placeholders.popLast() else {
+        logger.fault("Invalid placeholder in \(input)")
+        return input
+      }
+      guard let displayName = nameForSnippet(placeholderBody) else {
+        logger.fault("Failed to decode placeholder \(placeholderBody) in \(input)")
+        return input
+      }
+      let placeholder =
+        clientSupportsSnippets
+        ? formatLSPPlaceholder(displayName, number: number)
+        : ""
+      if placeholders.isEmpty {
+        result += placeholder
+      } else {
+        placeholders.latest.contents += placeholder
+      }
     }
-    let snippet = clientSupportsSnippets ? "${\(index):\(displayName)}" : ""
-    result.replaceSubrange(start.lowerBound..<end.upperBound, with: snippet)
-    index += 1
   }
+
   return result
 }
 
-/// Parse a SourceKit placeholder and extract the display name suitable for a
-/// LSP snippet.
-fileprivate func nameForSnippet(_ text: String) -> String? {
-  var text = text
+/// Scan `input` to identify special elements within: curly braces, which may
+/// need to be escaped; and SourceKit placeholder open/close delimiters.
+private func tokenize(_ input: String) -> [SnippetToken] {
+  var index = input.startIndex
+  var isAtEnd: Bool { index == input.endIndex }
+  func match(_ char: Character) -> Bool {
+    if isAtEnd || input[index] != char {
+      return false
+    } else {
+      input.formIndex(after: &index)
+      return true
+    }
+  }
+  func next() -> Character? {
+    guard !isAtEnd else { return nil }
+    defer { input.formIndex(after: &index) }
+    return input[index]
+  }
+
+  var tokens: [SnippetToken] = []
+  var text = ""
+  while let char = next() {
+    switch char {
+    case "<":
+      if match("#") {
+        tokens.append(.text(text))
+        text.removeAll()
+        tokens.append(.open)
+      } else {
+        text.append(char)
+      }
+
+    case "#":
+      if match(">") {
+        tokens.append(.text(text))
+        text.removeAll()
+        tokens.append(.close)
+      } else {
+        text.append(char)
+      }
+
+    case "{", "}":
+      tokens.append(.text(text))
+      text.removeAll()
+      tokens.append(.curlyBrace(char))
+
+    case let c:
+      text.append(c)
+    }
+  }
+
+  tokens.append(.text(text))
+
+  return tokens
+}
+
+/// A syntactical element inside a SourceKit snippet.
+private enum SnippetToken {
+  /// A placeholder delimiter.
+  case open, close
+  /// One of '{' or '}', which may need to be escaped in the output.
+  case curlyBrace(Character)
+  /// Any other consecutive run of characters from the input, which needs no
+  /// special treatment.
+  case text(String)
+}
+
+/// Given the interior text of a SourceKit placeholder, extract a display name
+/// suitable for a LSP snippet.
+private func nameForSnippet(_ body: String) -> String? {
+  var text = rewrappedAsPlaceholder(body)
   return text.withSyntaxText {
     guard let data = RawEditorPlaceholderData(syntaxText: $0) else {
       return nil
     }
     return String(syntaxText: data.typeForExpansionText ?? data.displayText)
   }
 }
+
+private let placeholderStart = "<#"
+private let placeholderEnd = "#>"
+private func rewrappedAsPlaceholder(_ body: String) -> String {
+  return placeholderStart + body + placeholderEnd
+}
+
+/// Wrap `body` in LSP snippet placeholder syntax, using `number` as the
+/// placeholder's index in the snippet.
+private func formatLSPPlaceholder(_ body: String, number: Int) -> String {
+  "${\(number):\(body)}"
+}
+
+private extension Array {
+  /// Mutable access to the final element of an array.
+  ///
+  /// - precondition: The array must not be empty.
+  var latest: Element {
+    get { self.last! }
+    _modify {
+      let index = self.index(before: self.endIndex)
+      yield &self[index]
+    }
+  }
+}

Tests/SourceKitLSPTests/RewriteSourceKitPlaceholdersTests.swift

@@ -0,0 +1,62 @@
+import SKTestSupport
+@testable import SourceKitLSP
+import XCTest
+
+final class RewriteSourceKitPlaceholdersTests : XCTestCase {
+  func testClientDoesNotSupportSnippets() {
+    let input = "foo(bar: \(placeholder: "T##Int##Int#"))"
+    let rewritten = rewriteSourceKitPlaceholders(in: input, clientSupportsSnippets: false)
+
+    XCTAssertEqual(rewritten, "foo(bar: )")
+  }
+
+  func testInputWithoutPlaceholders() {
+    let input = "foo()"
+    let rewritten = rewriteSourceKitPlaceholders(in: input, clientSupportsSnippets: true)
+
+    XCTAssertEqual(rewritten, "foo()")
+  }
+
+  func testPlaceholderWithType() {
+    let input = "foo(bar: \(placeholder: "T##bar##Int"))"
+    let rewritten = rewriteSourceKitPlaceholders(in: input, clientSupportsSnippets: true)
+
+    XCTAssertEqual(rewritten, "foo(bar: ${1:Int})")
+  }
+
+  func testMultiplePlaceholders() {
+    let input = "foo(bar: \(placeholder: "T##Int##Int"), baz: \(placeholder: "T##String##String"), quux: \(placeholder: "T##String##String"))"
+    let rewritten = rewriteSourceKitPlaceholders(in: input, clientSupportsSnippets: true)
+
+    XCTAssertEqual(rewritten, "foo(bar: ${1:Int}, baz: ${2:String}, quux: ${3:String})")
+  }
+
+  func testClosurePlaceholderReturnType() {
+    let input = "foo(bar: \(placeholder: "{ \(placeholder: "T##Int##Int") }"))"
+    let rewritten = rewriteSourceKitPlaceholders(in: input, clientSupportsSnippets: true)
+
+    XCTAssertEqual(rewritten, "foo(bar: ${1:\\{ ${2:Int} \\}})")
+  }
+
+  func testClosurePlaceholderArgumentType() {
+    let input = "foo(bar: \(placeholder: "{ \(placeholder: "T##Int##Int") in \(placeholder: "T##Void##Void") }"))"
+    let rewritten = rewriteSourceKitPlaceholders(in: input, clientSupportsSnippets: true)
+
+    XCTAssertEqual(rewritten, "foo(bar: ${1:\\{ ${2:Int} in ${3:Void} \\}})")
+  }
+
+  func testMultipleClosurePlaceholders() {
+    let input = "foo(\(placeholder: "{ \(placeholder: "T##Int##Int") }"), baz: \(placeholder: "{ \(placeholder: "Int") in \(placeholder: "T##Bool##Bool") }"))"
+    let rewritten = rewriteSourceKitPlaceholders(in: input, clientSupportsSnippets: true)
+
+    XCTAssertEqual(rewritten, "foo(${1:\\{ ${2:Int} \\}}, baz: ${3:\\{ ${4:Int} in ${5:Bool} \\}})")
+  }
+}
+
+private let placeholderStart = "<#"
+private let placeholderEnd = "#>"
+fileprivate extension DefaultStringInterpolation {
+  mutating func appendInterpolation(placeholder string: String) {
+    self.appendLiteral(placeholderStart + string + placeholderEnd)
+  }
+}

Tests/SourceKitLSPTests/SwiftCompletionTests.swift

@@ -864,18 +864,14 @@ final class SwiftCompletionTests: XCTestCase {
           sortText: nil,
           filterText: "myMap(:)",
           insertText: """
-            myMap { ${1:Int} in
-                ${2:Bool}
-            }
+            myMap(${1:\\{ ${2:Int} in ${3:Bool} \\}})
             """,
           insertTextFormat: .snippet,
           textEdit: .textEdit(
             TextEdit(
               range: Range(positions["1️⃣"]),
               newText: """
-                myMap { ${1:Int} in
-                    ${2:Bool}
-                }
+                myMap(${1:\\{ ${2:Int} in ${3:Bool} \\}})
                 """
             )
           )
@@ -912,18 +908,14 @@ final class SwiftCompletionTests: XCTestCase {
           sortText: nil,
           filterText: ".myMap(:)",
           insertText: """
-            ?.myMap { ${1:Int} in
-                ${2:Bool}
-            }
+            ?.myMap(${1:\\{ ${2:Int} in ${3:Bool} \\}})
             """,
           insertTextFormat: .snippet,
           textEdit: .textEdit(
             TextEdit(
               range: positions["1️⃣"]..<positions["2️⃣"],
               newText: """
-                ?.myMap { ${1:Int} in
-                    ${2:Bool}
-                }
+                ?.myMap(${1:\\{ ${2:Int} in ${3:Bool} \\}})
                 """
             )
           )
@@ -960,22 +952,14 @@ final class SwiftCompletionTests: XCTestCase {
           sortText: nil,
           filterText: "myMap(::)",
           insertText: """
-            myMap { ${1:Int} in
-                ${2:Bool}
-            } _: { ${3:Int} in
-                ${4:String}
-            }
+            myMap(${1:\\{ ${2:Int} in ${3:Bool} \\}}, ${4:\\{ ${5:Int} in ${6:String} \\}})
             """,
           insertTextFormat: .snippet,
           textEdit: .textEdit(
             TextEdit(
               range: Range(positions["1️⃣"]),
               newText: """
-                myMap { ${1:Int} in
-                    ${2:Bool}
-                } _: { ${3:Int} in
-                    ${4:String}
-                }
+                myMap(${1:\\{ ${2:Int} in ${3:Bool} \\}}, ${4:\\{ ${5:Int} in ${6:String} \\}})
                 """
             )
           )
@@ -1012,22 +996,14 @@ final class SwiftCompletionTests: XCTestCase {
           sortText: nil,
           filterText: "myMap(:second:)",
           insertText: """
-            myMap { ${1:Int} in
-                ${2:Bool}
-            } second: { ${3:Int} in
-                ${4:String}
-            }
+            myMap(${1:\\{ ${2:Int} in ${3:Bool} \\}}, second: ${4:\\{ ${5:Int} in ${6:String} \\}})
             """,
           insertTextFormat: .snippet,
           textEdit: .textEdit(
             TextEdit(
               range: Range(positions["1️⃣"]),
               newText: """
-                myMap { ${1:Int} in
-                    ${2:Bool}
-                } second: { ${3:Int} in
-                    ${4:String}
-                }
+                myMap(${1:\\{ ${2:Int} in ${3:Bool} \\}}, second: ${4:\\{ ${5:Int} in ${6:String} \\}})
                 """
             )
           )
@@ -1066,18 +1042,14 @@ final class SwiftCompletionTests: XCTestCase {
           sortText: nil,
           filterText: "myMap(:)",
           insertText: """
-            myMap { ${1:Int} in
-              ${2:Bool}
-            }
+            myMap(${1:\\{ ${2:Int} in ${3:Bool} \\}})
             """,
           insertTextFormat: .snippet,
           textEdit: .textEdit(
             TextEdit(
               range: Range(positions["1️⃣"]),
               newText: """
-                myMap { ${1:Int} in
-                  ${2:Bool}
-                }
+                myMap(${1:\\{ ${2:Int} in ${3:Bool} \\}})
                 """
             )
           )