Register diagnostics in currentDiagnostics when performing a diagnostic pull request

Sources/SourceKitLSP/Swift/SwiftLanguageServer.swift

@@ -271,7 +271,36 @@ public final class SwiftLanguageServer: ToolchainLanguageServer {
       }
     })
   }
-
+
+  /// Register the diagnostics returned from sourcekitd in `currentDiagnostics`
+  /// and returns the corresponding LSP diagnostics.
+  private func registerDiagnostics(
+    sourcekitdDiagnostics: SKDResponseArray?,
+    snapshot: DocumentSnapshot,
+    stage: DiagnosticStage
+  ) -> [Diagnostic] {
+    let supportsCodeDescription = capabilityRegistry.clientHasDiagnosticsCodeDescriptionSupport
+
+    var newDiags: [CachedDiagnostic] = []
+    sourcekitdDiagnostics?.forEach { _, diag in
+      if let diag = CachedDiagnostic(diag, in: snapshot, useEducationalNoteAsCode: supportsCodeDescription) {
+        newDiags.append(diag)
+      }
+      return true
+    }
+
+    let result = mergeDiagnostics(
+      old: currentDiagnostics[snapshot.document.uri] ?? [],
+      new: newDiags,
+      stage: stage,
+      isFallback: self.commandsByFile[snapshot.document.uri]?.isFallback ?? true
+    )
+    currentDiagnostics[snapshot.document.uri] = result
+
+    return result.map(\.diagnostic)
+
+  }
+
   /// Publish diagnostics for the given `snapshot`. We withhold semantic diagnostics if we are using
   /// fallback arguments.
   ///
@@ -287,31 +316,22 @@ public final class SwiftLanguageServer: ToolchainLanguageServer {
       return
     }
 
-    let isFallback = compileCommand?.isFallback ?? true
-
     let stageUID: sourcekitd_uid_t? = response[sourcekitd.keys.diagnostic_stage]
     let stage = stageUID.flatMap { DiagnosticStage($0, sourcekitd: sourcekitd) } ?? .sema
 
-    let supportsCodeDescription = capabilityRegistry.clientHasDiagnosticsCodeDescriptionSupport
-
-    // 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) {
-        newDiags.append(diag)
-      }
-      return true
-    }
-
-    let result = mergeDiagnostics(
-      old: currentDiagnostics[documentUri] ?? [],
-      new: newDiags, stage: stage, isFallback: isFallback)
-    currentDiagnostics[documentUri] = result
-
-    client.send(PublishDiagnosticsNotification(
-        uri: documentUri, version: snapshot.version, diagnostics: result.map { $0.diagnostic }))
+    let diagnostics = registerDiagnostics(
+      sourcekitdDiagnostics: response[keys.diagnostics],
+      snapshot: snapshot,
+      stage: stage
+    )
+
+    client.send(
+      PublishDiagnosticsNotification(
+        uri: documentUri,
+        version: snapshot.version,
+        diagnostics: diagnostics
+      )
+    )
   }
 
   /// Should be called on self.queue.
@@ -1363,20 +1383,16 @@ extension SwiftLanguageServer {
       skreq[keys.compilerargs] = compileCommand.compilerArgs
     }
 
-    let supportsCodeDescription = capabilityRegistry.clientHasDiagnosticsCodeDescriptionSupport
-
     let handle = self.sourcekitd.send(skreq, self.queue) { response in
       guard let dict = response.success else {
         return completion(.failure(ResponseError(response.failure!)))
       }
 
-      var diagnostics: [Diagnostic] = []
-      dict[keys.diagnostics]?.forEach { _, diag in
-        if let diagnostic = Diagnostic(diag, in: snapshot, useEducationalNoteAsCode: supportsCodeDescription) {
-          diagnostics.append(diagnostic)
-        }
-        return true
-      }
+      let diagnostics = self.registerDiagnostics(
+        sourcekitdDiagnostics: dict[keys.diagnostics],
+        snapshot: snapshot,
+        stage: .sema
+      )
 
       completion(.success(diagnostics))
     }

Tests/SourceKitLSPTests/PullDiagnosticsTests.swift

@@ -34,7 +34,10 @@ final class PullDiagnosticsTests: XCTestCase {
       rootPath: nil,
       rootURI: nil,
       initializationOptions: nil,
-      capabilities: ClientCapabilities(workspace: nil, textDocument: nil),
+      capabilities: ClientCapabilities(
+        workspace: nil,
+        textDocument: .init(codeAction: .init(codeActionLiteralSupport: .init(codeActionKind: .init(valueSet: [.quickFix]))))
+      ),
       trace: .off,
       workspaceFolders: nil
     ))
@@ -78,6 +81,43 @@ final class PullDiagnosticsTests: XCTestCase {
     }
     """)
     XCTAssertEqual(diagnostics.count, 1)
-    XCTAssertEqual(diagnostics[0].range, Position(line: 1, utf16index: 2)..<Position(line: 1, utf16index: 9))
+    let diagnostic = try XCTUnwrap(diagnostics.first)
+    XCTAssertEqual(diagnostic.range, Position(line: 1, utf16index: 2)..<Position(line: 1, utf16index: 9))
+  }
+
+  /// Test that we can get code actions for pulled diagnostics (https://github.com/apple/sourcekit-lsp/issues/776)
+  func testCodeActions() throws {
+    let diagnostics = try performDiagnosticRequest(text: """
+    protocol MyProtocol {
+      func bar()
+    }
+
+    struct Test: MyProtocol {}
+    """)
+    XCTAssertEqual(diagnostics.count, 1)
+    let diagnostic = try XCTUnwrap(diagnostics.first)
+    XCTAssertEqual(diagnostic.range, Position(line: 4, utf16index: 7)..<Position(line: 4, utf16index: 7))
+    let note = try XCTUnwrap(diagnostic.relatedInformation?.first)
+    XCTAssertEqual(note.location.range, Position(line: 4, utf16index: 7)..<Position(line: 4, utf16index: 7))
+    XCTAssertEqual(note.codeActions?.count ?? 0, 1)
+
+    let response = try sk.sendSync(CodeActionRequest(
+      range: note.location.range,
+      context: CodeActionContext(
+        diagnostics: diagnostics,
+        only: [.quickFix],
+        triggerKind: .invoked
+      ),
+      textDocument: TextDocumentIdentifier(note.location.uri)
+    ))
+
+    guard case .codeActions(let actions) = response else {
+      XCTFail("Expected codeActions response")
+      return
+    }
+
+    XCTAssertEqual(actions.count, 1)
+    let action = try XCTUnwrap(actions.first)
+    XCTAssertEqual(action.title, "do you want to add protocol stubs?")
   }
 }