Handle diagnostics in secondary files correctly

rdar://130503535

Author: ahoppen

Sources/SourceKitLSP/Swift/Diagnostic.swift

@@ -154,16 +154,27 @@ fileprivate extension String {
 extension Diagnostic {
 
   /// Creates a diagnostic from a sourcekitd response dictionary.
+  ///
+  /// `snapshot` is the snapshot of the document for which the diagnostics are generated.
+  /// `documentManager` is used to resolve positions of notes in secondary files.
   init?(
     _ diag: SKDResponseDictionary,
     in snapshot: DocumentSnapshot,
+    documentManager: DocumentManager,
     useEducationalNoteAsCode: Bool
   ) {
-    // FIXME: this assumes that the diagnostics are all in the same file.
-
     let keys = diag.sourcekitd.keys
     let values = diag.sourcekitd.values
 
+    guard let filePath: String = diag[keys.filePath] else {
+      logger.fault("Missing file path in diagnostic")
+      return nil
+    }
+    guard filePath == snapshot.uri.pseudoPath else {
+      logger.error("Ignoring diagnostic from a different file: \(filePath)")
+      return nil
+    }
+
     guard let message: String = diag[keys.description]?.withFirstLetterUppercased() else { return nil }
 
     var range: Range<Position>? = nil
@@ -237,7 +248,13 @@ 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,
+            primaryDocumentSnapshot: snapshot,
+            documentManager: documentManager
+          )
+        else { return true }
         notes?.append(note)
         return true
       }
@@ -309,9 +326,28 @@ extension Diagnostic {
 extension DiagnosticRelatedInformation {
 
   /// Creates related information from a sourcekitd note response dictionary.
-  init?(_ diag: SKDResponseDictionary, in snapshot: DocumentSnapshot) {
+  ///
+  /// `primaryDocumentSnapshot` is the snapshot of the document for which the diagnostics are generated.
+  /// `documentManager` is used to resolve positions of notes in secondary files.
+  init?(_ diag: SKDResponseDictionary, primaryDocumentSnapshot: DocumentSnapshot, documentManager: DocumentManager) {
     let keys = diag.sourcekitd.keys
 
+    guard let filePath: String = diag[keys.filePath] else {
+      logger.fault("Missing file path in related diagnostic information")
+      return nil
+    }
+    let uri = DocumentURI(filePath: filePath, isDirectory: false)
+    let snapshot: DocumentSnapshot
+    if filePath == primaryDocumentSnapshot.uri.pseudoPath {
+      snapshot = primaryDocumentSnapshot
+    } else if let inMemorySnapshot = try? documentManager.latestSnapshot(uri) {
+      snapshot = inMemorySnapshot
+    } else if let snapshotFromDisk = try? DocumentSnapshot(withContentsFromDisk: uri, language: .swift) {
+      snapshot = snapshotFromDisk
+    } else {
+      return nil
+    }
+
     var position: Position? = nil
     if let line: Int = diag[keys.line],
       let utf8Column: Int = diag[keys.column],

Sources/SourceKitLSP/Swift/DiagnosticReportManager.swift

@@ -113,6 +113,7 @@ actor DiagnosticReportManager {
         Diagnostic(
           diag,
           in: snapshot,
+          documentManager: documentManager,
           useEducationalNoteAsCode: self.clientHasDiagnosticsCodeDescriptionSupport
         )
       }) ?? []

Tests/SourceKitLSPTests/PullDiagnosticsTests.swift

@@ -349,4 +349,30 @@ final class PullDiagnosticsTests: XCTestCase {
     diagnosticRequestCancelled.fulfill()
     try await fulfillmentOfOrThrow([diagnosticResponseReceived])
   }
+
+  func testNoteInSecondaryFile() async throws {
+    let project = try await SwiftPMTestProject(files: [
+      "FileA.swift": """
+      @available(*, unavailable)
+      struct 1️⃣Test {}
+      """,
+      "FileB.swift": """
+      func test() {
+          _ = Test()
+      }
+      """,
+    ])
+
+    let (uri, _) = try project.openDocument("FileB.swift")
+    let diagnostics = try await project.testClient.send(
+      DocumentDiagnosticsRequest(textDocument: TextDocumentIdentifier(uri))
+    )
+    guard case .full(let diagnostics) = diagnostics else {
+      XCTFail("Expected full diagnostics report")
+      return
+    }
+    let diagnostic = try XCTUnwrap(diagnostics.items.only)
+    let note = try XCTUnwrap(diagnostic.relatedInformation?.only)
+    XCTAssertEqual(note.location, try project.location(from: "1️⃣", to: "1️⃣", in: "FileA.swift"))
+  }
 }