From 8ba7262f21091164381adf85f492d9969448e106 Mon Sep 17 00:00:00 2001
From: Alex Hoppen <ahoppen@apple.com>
Date: Tue, 26 Apr 2022 14:30:01 +0200
Subject: [PATCH] Translate placeholders to LSP snippets in CodeActions

Previously, we would insert `<#code#>` placeholders when inserting missing protocol requirements (both via the Fix-It and the refactor command). Translate them to LSP placeholders before inserting them.

rdar://92161144
---
 .../SKTestSupport/INPUTS/Fixit/Fixit.swift    |  7 ++
 .../SKTestSupport/INPUTS/Fixit/project.json   |  1 +
 .../SourceKitLSP/Swift/CodeCompletion.swift   | 22 +----
 Sources/SourceKitLSP/Swift/Diagnostic.swift   | 26 +++---
 .../Swift/EditorPlaceholder.swift             | 22 +++++
 .../Swift/SemanticRefactoring.swift           |  8 +-
 .../Swift/SwiftLanguageServer.swift           |  4 +-
 Tests/SourceKitLSPTests/CodeActionTests.swift | 85 +++++++++++++++++++
 8 files changed, 139 insertions(+), 36 deletions(-)
 create mode 100644 Sources/SKTestSupport/INPUTS/Fixit/Fixit.swift
 create mode 100644 Sources/SKTestSupport/INPUTS/Fixit/project.json

diff --git a/Sources/SKTestSupport/INPUTS/Fixit/Fixit.swift b/Sources/SKTestSupport/INPUTS/Fixit/Fixit.swift
new file mode 100644
index 00000000..b2f0a004
--- /dev/null
+++ b/Sources/SKTestSupport/INPUTS/Fixit/Fixit.swift
@@ -0,0 +1,7 @@
+protocol MyProto {
+  func foo()
+}
+
+struct /*MyStruct:def*/MyStruct: MyProto {
+
+}
diff --git a/Sources/SKTestSupport/INPUTS/Fixit/project.json b/Sources/SKTestSupport/INPUTS/Fixit/project.json
new file mode 100644
index 00000000..25de8a7c
--- /dev/null
+++ b/Sources/SKTestSupport/INPUTS/Fixit/project.json
@@ -0,0 +1 @@
+{ "sources": ["Fixit.swift"] }
diff --git a/Sources/SourceKitLSP/Swift/CodeCompletion.swift b/Sources/SourceKitLSP/Swift/CodeCompletion.swift
index 815ff432..1cc6c738 100644
--- a/Sources/SourceKitLSP/Swift/CodeCompletion.swift
+++ b/Sources/SourceKitLSP/Swift/CodeCompletion.swift
@@ -142,7 +142,7 @@ extension SwiftLanguageServer {
       let clientCompletionCapabilities = self.clientCapabilities.textDocument?.completion
       let clientSupportsSnippets = clientCompletionCapabilities?.completionItem?.snippetSupport == true
       let text = insertText.map {
-        self.rewriteSourceKitPlaceholders(inString: $0, clientSupportsSnippets: clientSupportsSnippets)
+        rewriteSourceKitPlaceholders(inString: $0, clientSupportsSnippets: clientSupportsSnippets)
       }
       let isInsertTextSnippet = clientSupportsSnippets && text != insertText
 
@@ -189,26 +189,6 @@ extension SwiftLanguageServer {
     return result
   }
 
-  func rewriteSourceKitPlaceholders(inString string: String, clientSupportsSnippets: Bool) -> String {
-    var result = string
-    var index = 1
-    while let start = result.range(of: EditorPlaceholder.placeholderPrefix) {
-      guard let end = result[start.upperBound...].range(of: EditorPlaceholder.placeholderSuffix) else {
-        log("invalid placeholder in \(string)", level: .debug)
-        return string
-      }
-      let rawPlaceholder = String(result[start.lowerBound..<end.upperBound])
-      guard let displayName = EditorPlaceholder(rawPlaceholder)?.displayName else {
-        log("failed to decode placeholder \(rawPlaceholder) in \(string)", level: .debug)
-        return string
-      }
-      let placeholder = clientSupportsSnippets ? "${\(index):\(displayName)}" : ""
-      result.replaceSubrange(start.lowerBound..<end.upperBound, with: placeholder)
-      index += 1
-    }
-    return result
-  }
-
   /// Adjust completion position to the start of identifier characters.
   func adjustCompletionLocation(_ pos: Position, in snapshot: DocumentSnapshot) -> Position? {
     guard pos.line < snapshot.lineTable.count else {
diff --git a/Sources/SourceKitLSP/Swift/Diagnostic.swift b/Sources/SourceKitLSP/Swift/Diagnostic.swift
index 94b6f178..a095bdf4 100644
--- a/Sources/SourceKitLSP/Swift/Diagnostic.swift
+++ b/Sources/SourceKitLSP/Swift/Diagnostic.swift
@@ -20,10 +20,10 @@ extension CodeAction {
   /// Creates a CodeAction from a list for sourcekit fixits.
   ///
   /// If this is from a note, the note's description should be passed as `fromNote`.
-  init?(fixits: SKDResponseArray, in snapshot: DocumentSnapshot, fromNote: String?) {
+  init?(fixits: SKDResponseArray, in snapshot: DocumentSnapshot, fromNote: String?, clientSupportsSnippets: Bool) {
     var edits: [TextEdit] = []
     let editsMapped = fixits.forEach { (_, skfixit) -> Bool in
-      if let edit = TextEdit(fixit: skfixit, in: snapshot) {
+      if let edit = TextEdit(fixit: skfixit, in: snapshot, clientSupportsSnippets: clientSupportsSnippets) {
         edits.append(edit)
         return true
       }
@@ -88,7 +88,7 @@ extension CodeAction {
 extension TextEdit {
 
   /// Creates a TextEdit from a sourcekitd fixit response dictionary.
-  init?(fixit: SKDResponseDictionary, in snapshot: DocumentSnapshot) {
+  init?(fixit: SKDResponseDictionary, in snapshot: DocumentSnapshot, clientSupportsSnippets: Bool) {
     let keys = fixit.sourcekitd.keys
     if let utf8Offset: Int = fixit[keys.offset],
        let length: Int = fixit[keys.length],
@@ -97,7 +97,8 @@ extension TextEdit {
        let endPosition = snapshot.positionOf(utf8Offset: utf8Offset + length),
        length > 0 || !replacement.isEmpty
     {
-      self.init(range: position..<endPosition, newText: replacement)
+      let replacementWithSnippets = rewriteSourceKitPlaceholders(inString: replacement, clientSupportsSnippets: clientSupportsSnippets)
+      self.init(range: position..<endPosition, newText: replacementWithSnippets)
     } else {
       return nil
     }
@@ -109,7 +110,8 @@ extension Diagnostic {
   /// Creates a diagnostic from a sourcekitd response dictionary.
   init?(_ diag: SKDResponseDictionary,
         in snapshot: DocumentSnapshot,
-        useEducationalNoteAsCode: Bool) {
+        useEducationalNoteAsCode: Bool,
+        clientSupportsSnippets: Bool) {
     // FIXME: this assumes that the diagnostics are all in the same file.
 
     let keys = diag.sourcekitd.keys
@@ -177,7 +179,7 @@ extension Diagnostic {
 
     var actions: [CodeAction]? = nil
     if let skfixits: SKDResponseArray = diag[keys.fixits],
-       let action = CodeAction(fixits: skfixits, in: snapshot, fromNote: nil) {
+       let action = CodeAction(fixits: skfixits, in: snapshot, fromNote: nil, clientSupportsSnippets: clientSupportsSnippets) {
       actions = [action]
     }
 
@@ -185,7 +187,7 @@ extension Diagnostic {
     if let sknotes: SKDResponseArray = diag[keys.diagnostics] {
       notes = []
       sknotes.forEach { (_, sknote) -> Bool in
-        guard let note = DiagnosticRelatedInformation(sknote, in: snapshot) else { return true }
+        guard let note = DiagnosticRelatedInformation(sknote, in: snapshot, clientSupportsSnippets: clientSupportsSnippets) else { return true }
         notes?.append(note)
         return true
       }
@@ -222,7 +224,7 @@ extension Diagnostic {
 extension DiagnosticRelatedInformation {
 
   /// Creates related information from a sourcekitd note response dictionary.
-  init?(_ diag: SKDResponseDictionary, in snapshot: DocumentSnapshot) {
+  init?(_ diag: SKDResponseDictionary, in snapshot: DocumentSnapshot, clientSupportsSnippets: Bool) {
     let keys = diag.sourcekitd.keys
 
     var position: Position? = nil
@@ -243,7 +245,7 @@ extension DiagnosticRelatedInformation {
 
     var actions: [CodeAction]? = nil
     if let skfixits: SKDResponseArray = diag[keys.fixits],
-       let action = CodeAction(fixits: skfixits, in: snapshot, fromNote: message) {
+       let action = CodeAction(fixits: skfixits, in: snapshot, fromNote: message, clientSupportsSnippets: clientSupportsSnippets) {
       actions = [action]
     }
 
@@ -277,11 +279,13 @@ struct CachedDiagnostic {
 extension CachedDiagnostic {
   init?(_ diag: SKDResponseDictionary,
         in snapshot: DocumentSnapshot,
-        useEducationalNoteAsCode: Bool) {
+        useEducationalNoteAsCode: Bool,
+        clientSupportsSnippets: Bool) {
     let sk = diag.sourcekitd
     guard let diagnostic = Diagnostic(diag,
                                       in: snapshot,
-                                      useEducationalNoteAsCode: useEducationalNoteAsCode) else {
+                                      useEducationalNoteAsCode: useEducationalNoteAsCode,
+                                      clientSupportsSnippets: clientSupportsSnippets) else {
       return nil
     }
     self.diagnostic = diagnostic
diff --git a/Sources/SourceKitLSP/Swift/EditorPlaceholder.swift b/Sources/SourceKitLSP/Swift/EditorPlaceholder.swift
index ad04e074..d922ae64 100644
--- a/Sources/SourceKitLSP/Swift/EditorPlaceholder.swift
+++ b/Sources/SourceKitLSP/Swift/EditorPlaceholder.swift
@@ -10,6 +10,8 @@
 //
 //===----------------------------------------------------------------------===//
 
+import LSPLogging
+
 public enum EditorPlaceholder: Hashable {
   case basic(String)
   case typed(displayName: String, type: String, typeForExpansion: String)
@@ -52,3 +54,23 @@ public enum EditorPlaceholder: Hashable {
     }
   }
 }
+
+func rewriteSourceKitPlaceholders(inString string: String, clientSupportsSnippets: Bool) -> String {
+  var result = string
+  var index = 1
+  while let start = result.range(of: EditorPlaceholder.placeholderPrefix) {
+    guard let end = result[start.upperBound...].range(of: EditorPlaceholder.placeholderSuffix) else {
+      log("invalid placeholder in \(string)", level: .debug)
+      return string
+    }
+    let rawPlaceholder = String(result[start.lowerBound..<end.upperBound])
+    guard let displayName = EditorPlaceholder(rawPlaceholder)?.displayName else {
+      log("failed to decode placeholder \(rawPlaceholder) in \(string)", level: .debug)
+      return string
+    }
+    let placeholder = clientSupportsSnippets ? "${\(index):\(displayName)}" : ""
+    result.replaceSubrange(start.lowerBound..<end.upperBound, with: placeholder)
+    index += 1
+  }
+  return result
+}
diff --git a/Sources/SourceKitLSP/Swift/SemanticRefactoring.swift b/Sources/SourceKitLSP/Swift/SemanticRefactoring.swift
index 6fc44c0c..02e10dfe 100644
--- a/Sources/SourceKitLSP/Swift/SemanticRefactoring.swift
+++ b/Sources/SourceKitLSP/Swift/SemanticRefactoring.swift
@@ -36,7 +36,7 @@ struct SemanticRefactoring {
   ///   - dict: Response dictionary to extract information from.
   ///   - url: The client URL that triggered the `semantic_refactoring` request.
   ///   - keys: The sourcekitd key set to use for looking up into `dict`.
-  init?(_ title: String, _ dict: SKDResponseDictionary, _ snapshot: DocumentSnapshot, _ keys: sourcekitd_keys) {
+  init?(_ title: String, _ dict: SKDResponseDictionary, _ snapshot: DocumentSnapshot, _ keys: sourcekitd_keys, clientSupportsSnippets: Bool) {
     guard let categorizedEdits: SKDResponseArray = dict[keys.categorizededits] else {
       return nil
     }
@@ -59,7 +59,8 @@ struct SemanticRefactoring {
                                                  utf8Column: endColumn - 1),
            let text: String = value[keys.text]
         {
-          let edit = TextEdit(range: startPosition..<endPosition, newText: text)
+          let textWithSnippets = rewriteSourceKitPlaceholders(inString: text, clientSupportsSnippets: clientSupportsSnippets)
+          let edit = TextEdit(range: startPosition..<endPosition, newText: textWithSnippets)
           textEdits.append(edit)
         }
         return true
@@ -164,7 +165,8 @@ extension SwiftLanguageServer {
         guard let dict = result.success else {
           return completion(.failure(.responseError(ResponseError(result.failure!))))
         }
-        guard let refactor = SemanticRefactoring(refactorCommand.title, dict, snapshot, self.keys) else {
+        let supportsSnippets = (self.clientCapabilities.textDocument?.completion?.completionItem?.snippetSupport == true)
+        guard let refactor = SemanticRefactoring(refactorCommand.title, dict, snapshot, self.keys, clientSupportsSnippets: supportsSnippets) else {
           return completion(.failure(.noEditsNeeded(uri)))
         }
         completion(.success(refactor))
diff --git a/Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift b/Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift
index 32a3a7bf..ab77c7d2 100644
--- a/Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift
+++ b/Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift
@@ -293,13 +293,15 @@ public final class SwiftLanguageServer: ToolchainLanguageServer {
 
     let supportsCodeDescription =
            (clientCapabilities.textDocument?.publishDiagnostics?.codeDescriptionSupport == true)
+    let supportsSnippets = (self.clientCapabilities.textDocument?.completion?.completionItem?.snippetSupport == true)
 
     // Note: we make the notification even if there are no diagnostics to clear the current state.
     var newDiags: [CachedDiagnostic] = []
     response[keys.diagnostics]?.forEach { _, diag in
       if let diag = CachedDiagnostic(diag,
                                      in: snapshot,
-                                     useEducationalNoteAsCode: supportsCodeDescription) {
+                                     useEducationalNoteAsCode: supportsCodeDescription,
+                                     clientSupportsSnippets: supportsSnippets) {
         newDiags.append(diag)
       }
       return true
diff --git a/Tests/SourceKitLSPTests/CodeActionTests.swift b/Tests/SourceKitLSPTests/CodeActionTests.swift
index bc58bc7b..90659d43 100644
--- a/Tests/SourceKitLSPTests/CodeActionTests.swift
+++ b/Tests/SourceKitLSPTests/CodeActionTests.swift
@@ -28,6 +28,7 @@ final class CodeActionTests: XCTestCase {
     let codeActionLiteralSupport = CodeActionLiteralSupport(codeActionKind: codeActionKinds)
     codeActionCapabilities.codeActionLiteralSupport = codeActionLiteralSupport
     documentCapabilities.codeAction = codeActionCapabilities
+    documentCapabilities.completion = .init(completionItem: .init(snippetSupport: true))
     return ClientCapabilities(workspace: nil, textDocument: documentCapabilities)
   }
 
@@ -246,4 +247,88 @@ final class CodeActionTests: XCTestCase {
 
     XCTAssertEqual(result, .codeActions([expectedCodeAction]))
   }
+
+  func testCodeActionsUseLSPSnippets() throws {
+    let capabilities = clientCapabilitiesWithCodeActionSupport()
+    let ws = try staticSourceKitTibsWorkspace(name: "Fixit", clientCapabilities: capabilities)!
+
+    let def = ws.testLoc("MyStruct:def")
+
+    try ws.openDocument(def.url, language: .swift)
+
+    let syntacticDiagnosticsReceived = self.expectation(description: "Syntactic diagnotistics received")
+    let semanticDiagnosticsReceived = self.expectation(description: "Semantic diagnotistics received")
+
+    ws.sk.appendOneShotNotificationHandler  { (note: Notification<PublishDiagnosticsNotification>) in
+      // syntactic diagnostics
+      XCTAssertEqual(note.params.uri, def.docUri)
+      XCTAssertEqual(note.params.diagnostics, [])
+      syntacticDiagnosticsReceived.fulfill()
+    }
+
+    var diags: [Diagnostic]! = nil
+    ws.sk.appendOneShotNotificationHandler  { (note: Notification<PublishDiagnosticsNotification>) in
+      // semantic diagnostics
+      XCTAssertEqual(note.params.uri, def.docUri)
+      XCTAssertEqual(note.params.diagnostics.count, 1)
+      diags = note.params.diagnostics
+      semanticDiagnosticsReceived.fulfill()
+    }
+
+    self.wait(for: [syntacticDiagnosticsReceived, semanticDiagnosticsReceived], timeout: 30)
+
+    let textDocument = TextDocumentIdentifier(def.url)
+    let actionsRequest = CodeActionRequest(range: def.position..<def.position, context: .init(diagnostics: diags), textDocument: textDocument)
+    let actionResult = try ws.sk.sendSync(actionsRequest)
+
+    guard case .codeActions(let codeActions) = actionResult else {
+      return XCTFail("Expected code actions, not commands as a response")
+    }
+
+    // Check that the Fix-It action contains snippets
+
+    guard let quickFixAction = codeActions.filter({ $0.kind == .quickFix }).spm_only else {
+      return XCTFail("Expected exactly one quick fix action")
+    }
+    guard let change = quickFixAction.edit?.changes?[def.docUri]?.spm_only else {
+      return XCTFail("Expected exactly one change")
+    }
+    XCTAssertEqual(change.newText, """
+
+        func foo() {
+            ${1:code}
+        }
+
+    """)
+
+    // Check that the refactor action contains snippets
+    guard let refactorAction = codeActions.filter({ $0.kind == .refactor }).spm_only else {
+      return XCTFail("Expected exactly one refactor action")
+    }
+    guard let command = refactorAction.command else {
+      return XCTFail("Expected the refactor action to have a command")
+    }
+
+    let editReceived = self.expectation(description: "Received ApplyEdit request")
+
+    ws.sk.appendOneShotRequestHandler { (request: Request<ApplyEditRequest>) in
+      defer {
+        editReceived.fulfill()
+      }
+      guard let change = request.params.edit.changes?[def.docUri]?.spm_only else {
+        return XCTFail("Expected exactly one edit")
+      }
+      XCTAssertEqual(change.newText, """
+
+          func foo() {
+              ${1:code}
+          }
+
+      """)
+      request.reply(ApplyEditResponse(applied: true, failureReason: nil))
+    }
+    _ = try ws.sk.sendSync(ExecuteCommandRequest(command: command.command, arguments: command.arguments))
+
+    self.wait(for: [editReceived], timeout: 5)
+  }
 }