Wait for target to be prepared before returning diagnostics

Diagnostics are usually not helpful if the file’s target hasn’t been prepared yet. We should await the target’s preparation before returning diagnostics.

rdar://128645617

Author: ahoppen

Package.swift

@@ -367,6 +367,7 @@ let package = Package(
       name: "SourceKitLSPTests",
       dependencies: [
         "BuildServerProtocol",
+        "CAtomics",
         "LSPLogging",
         "LSPTestSupport",
         "LanguageServerProtocol",

Sources/SKTestSupport/MultiFileTestProject.swift

@@ -82,6 +82,7 @@ public class MultiFileTestProject {
     workspaces: (URL) async throws -> [WorkspaceFolder] = { [WorkspaceFolder(uri: DocumentURI($0))] },
     capabilities: ClientCapabilities = ClientCapabilities(),
     serverOptions: SourceKitLSPServer.Options = .testDefault,
+    enableBackgroundIndexing: Bool = false,
     usePullDiagnostics: Bool = true,
     preInitialization: ((TestSourceKitLSPClient) -> Void)? = nil,
     cleanUp: (() -> Void)? = nil,
@@ -117,6 +118,7 @@ public class MultiFileTestProject {
       serverOptions: serverOptions,
       capabilities: capabilities,
       usePullDiagnostics: usePullDiagnostics,
+      enableBackgroundIndexing: enableBackgroundIndexing,
       workspaceFolders: workspaces(scratchDirectory),
       preInitialization: preInitialization,
       cleanUp: { [scratchDirectory] in

Sources/SKTestSupport/SwiftPMTestProject.swift

@@ -44,9 +44,10 @@ public class SwiftPMTestProject: MultiFileTestProject {
     allowBuildFailure: Bool = false,
     capabilities: ClientCapabilities = ClientCapabilities(),
     serverOptions: SourceKitLSPServer.Options = .testDefault,
+    enableBackgroundIndexing: Bool = false,
+    usePullDiagnostics: Bool = true,
     pollIndex: Bool = true,
     preInitialization: ((TestSourceKitLSPClient) -> Void)? = nil,
-    usePullDiagnostics: Bool = true,
     cleanUp: (() -> Void)? = nil,
     testName: String = #function
   ) async throws {
@@ -71,6 +72,7 @@ public class SwiftPMTestProject: MultiFileTestProject {
       workspaces: workspaces,
       capabilities: capabilities,
       serverOptions: serverOptions,
+      enableBackgroundIndexing: enableBackgroundIndexing,
       usePullDiagnostics: usePullDiagnostics,
       preInitialization: preInitialization,
       cleanUp: cleanUp,

Sources/SKTestSupport/TestSourceKitLSPClient.swift

@@ -88,7 +88,8 @@ public final class TestSourceKitLSPClient: MessageHandler {
   ///     `true` by default
   ///   - initializationOptions: Initialization options to pass to the SourceKit-LSP server.
   ///   - capabilities: The test client's capabilities.
-  ///   - usePullDiagnostics: Whether to use push diagnostics or use push-based diagnostics
+  ///   - usePullDiagnostics: Whether to use push diagnostics or use push-based diagnostics.
+  ///   - enableBackgroundIndexing: Whether background indexing should be enabled in the project.
   ///   - workspaceFolders: Workspace folders to open.
   ///   - preInitialization: A closure that is called after the test client is created but before SourceKit-LSP is
   ///     initialized. This can be used to eg. register request handlers.
@@ -102,6 +103,7 @@ public final class TestSourceKitLSPClient: MessageHandler {
     initializationOptions: LSPAny? = nil,
     capabilities: ClientCapabilities = ClientCapabilities(),
     usePullDiagnostics: Bool = true,
+    enableBackgroundIndexing: Bool = false,
     workspaceFolders: [WorkspaceFolder]? = nil,
     preInitialization: ((TestSourceKitLSPClient) -> Void)? = nil,
     cleanUp: @Sendable @escaping () -> Void = {}
@@ -115,6 +117,7 @@ public final class TestSourceKitLSPClient: MessageHandler {
     if let moduleCache {
       serverOptions.buildSetup.flags.swiftCompilerFlags += ["-module-cache-path", moduleCache.path]
     }
+    serverOptions.indexOptions.enableBackgroundIndexing = enableBackgroundIndexing
 
     var notificationYielder: AsyncStream<any NotificationType>.Continuation!
     self.notifications = AsyncStream { continuation in
@@ -155,8 +158,8 @@ public final class TestSourceKitLSPClient: MessageHandler {
         XCTAssertEqual(request.registrations.only?.method, DocumentDiagnosticsRequest.method)
         return VoidResponse()
       }
-      preInitialization?(self)
     }
+    preInitialization?(self)
     if initialize {
       _ = try await self.send(
         InitializeRequest(
@@ -193,12 +196,22 @@ public final class TestSourceKitLSPClient: MessageHandler {
   /// Send the request to `server` and return the request result.
   public func send<R: RequestType>(_ request: R) async throws -> R.Response {
     return try await withCheckedThrowingContinuation { continuation in
-      server.handle(request, id: .number(Int(nextRequestID.fetchAndIncrement()))) { result in
+      self.send(request) { result in
         continuation.resume(with: result)
       }
     }
   }
 
+  /// Send the request to `server` and return the result via a completion handler.
+  ///
+  /// This version of the `send` function should only be used if some action needs to be performed after the request is
+  /// sent but before it returns a result.
+  public func send<R: RequestType>(_ request: R, completionHandler: @escaping (LSPResult<R.Response>) -> Void) {
+    server.handle(request, id: .number(Int(nextRequestID.fetchAndIncrement()))) { result in
+      completionHandler(result)
+    }
+  }
+
   /// Send the notification to `server`.
   public func send(_ notification: some NotificationType) {
     server.handle(notification)

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -327,10 +327,7 @@ public final actor SemanticIndexManager {
   /// Schedule preparation of the target that contains the given URI, building all modules that the file depends on.
   ///
   /// This is intended to be called when the user is interacting with the document at the given URI.
-  public func schedulePreparationForEditorFunctionality(
-    of uri: DocumentURI,
-    priority: TaskPriority? = nil
-  ) {
+  public func schedulePreparationForEditorFunctionality(of uri: DocumentURI, priority: TaskPriority? = nil) {
     if inProgressPrepareForEditorTask?.document == uri {
       // We are already preparing this document, so nothing to do. This is necessary to avoid the following scenario:
       // Determining the canonical configured target for a document takes 1s and we get a new document request for the
@@ -341,13 +338,7 @@ public final actor SemanticIndexManager {
     let id = UUID()
     let task = Task(priority: priority) {
       await withLoggingScope("preparation") {
-        guard let target = await buildSystemManager.canonicalConfiguredTarget(for: uri) else {
-          return
-        }
-        if Task.isCancelled {
-          return
-        }
-        await self.prepare(targets: [target], priority: priority)
+        await self.prepareFileForEditorFunctionality(uri)
         if inProgressPrepareForEditorTask?.id == id {
           inProgressPrepareForEditorTask = nil
         }
@@ -357,6 +348,20 @@ public final actor SemanticIndexManager {
     inProgressPrepareForEditorTask = (id, uri, task)
   }
 
+  /// Prepare the target that the given file is in, building all modules that the file depends on. Returns when
+  /// preparation has finished.
+  ///
+  /// If file's target is known to be up-to-date, this returns almost immediately.
+  public func prepareFileForEditorFunctionality(_ uri: DocumentURI) async {
+    guard let target = await buildSystemManager.canonicalConfiguredTarget(for: uri) else {
+      return
+    }
+    if Task.isCancelled {
+      return
+    }
+    await self.prepare(targets: [target], priority: nil)
+  }
+
   // MARK: - Helper functions
 
   /// Prepare the given targets for indexing.

Sources/SourceKitLSP/Swift/SwiftLanguageService.swift

@@ -17,6 +17,7 @@ import LSPLogging
 import LanguageServerProtocol
 import SKCore
 import SKSupport
+import SemanticIndex
 import SourceKitD
 import SwiftParser
 import SwiftParserDiagnostics
@@ -123,6 +124,9 @@ public actor SwiftLanguageService: LanguageService, Sendable {
 
   let syntaxTreeManager = SyntaxTreeManager()
 
+  /// The `semanticIndexManager` of the workspace this language service was created for.
+  private let semanticIndexManager: SemanticIndexManager?
+
   nonisolated var keys: sourcekitd_api_keys { return sourcekitd.keys }
   nonisolated var requests: sourcekitd_api_requests { return sourcekitd.requests }
   nonisolated var values: sourcekitd_api_values { return sourcekitd.values }
@@ -192,6 +196,7 @@ public actor SwiftLanguageService: LanguageService, Sendable {
     self.swiftFormat = toolchain.swiftFormat
     self.sourcekitd = try await DynamicallyLoadedSourceKitD.getOrCreate(dylibPath: sourcekitd)
     self.capabilityRegistry = workspace.capabilityRegistry
+    self.semanticIndexManager = workspace.semanticIndexManager
     self.serverOptions = options
     self.documentManager = DocumentManager()
     self.state = .connected
@@ -875,6 +880,7 @@ extension SwiftLanguageService {
 
   public func documentDiagnostic(_ req: DocumentDiagnosticsRequest) async throws -> DocumentDiagnosticReport {
     do {
+      await semanticIndexManager?.prepareFileForEditorFunctionality(req.textDocument.uri)
       let snapshot = try documentManager.latestSnapshot(req.textDocument.uri)
       let buildSettings = await self.buildSettings(for: req.textDocument.uri)
       let diagnosticReport = try await self.diagnosticReportManager.diagnosticReport(

Tests/SourceKitLSPTests/BackgroundIndexingTests.swift

@@ -18,10 +18,6 @@ import SemanticIndex
 import SourceKitLSP
 import XCTest
 
-fileprivate let backgroundIndexingOptions = SourceKitLSPServer.Options(
-  indexOptions: IndexOptions(enableBackgroundIndexing: true)
-)
-
 final class BackgroundIndexingTests: XCTestCase {
   func testBackgroundIndexingOfSingleFile() async throws {
     let project = try await SwiftPMTestProject(
@@ -33,7 +29,7 @@ final class BackgroundIndexingTests: XCTestCase {
         }
         """
       ],
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
 
     let (uri, positions) = try project.openDocument("MyFile.swift")
@@ -76,7 +72,7 @@ final class BackgroundIndexingTests: XCTestCase {
         }
         """,
       ],
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
 
     let (uri, positions) = try project.openDocument("MyFile.swift")
@@ -134,7 +130,7 @@ final class BackgroundIndexingTests: XCTestCase {
           ]
         )
         """,
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
 
     let (uri, positions) = try project.openDocument("MyFile.swift")
@@ -166,7 +162,7 @@ final class BackgroundIndexingTests: XCTestCase {
   }
 
   func testBackgroundIndexingHappensWithLowPriority() async throws {
-    var serverOptions = backgroundIndexingOptions
+    var serverOptions = SourceKitLSPServer.Options.testDefault
     serverOptions.indexTestHooks.preparationTaskDidFinish = { taskDescription in
       XCTAssert(Task.currentPriority == .low, "\(taskDescription) ran with priority \(Task.currentPriority)")
     }
@@ -199,6 +195,7 @@ final class BackgroundIndexingTests: XCTestCase {
         )
         """,
       serverOptions: serverOptions,
+      enableBackgroundIndexing: true,
       pollIndex: false
     )
 
@@ -249,7 +246,7 @@ final class BackgroundIndexingTests: XCTestCase {
           ]
         )
         """,
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
 
     let dependencyUrl = try XCTUnwrap(
@@ -300,7 +297,7 @@ final class BackgroundIndexingTests: XCTestCase {
         }
         """,
       ],
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
 
     let (uri, positions) = try project.openDocument("MyFile.c")
@@ -339,7 +336,7 @@ final class BackgroundIndexingTests: XCTestCase {
     let receivedReportProgressNotification = self.expectation(
       description: "Received work done progress saying indexing"
     )
-    var serverOptions = backgroundIndexingOptions
+    var serverOptions = SourceKitLSPServer.Options.testDefault
     serverOptions.indexTestHooks = IndexTestHooks(
       buildGraphGenerationDidFinish: {
         await self.fulfillment(of: [receivedBeginProgressNotification], timeout: defaultTimeout)
@@ -359,6 +356,7 @@ final class BackgroundIndexingTests: XCTestCase {
       ],
       capabilities: ClientCapabilities(window: WindowClientCapabilities(workDoneProgress: true)),
       serverOptions: serverOptions,
+      enableBackgroundIndexing: true,
       pollIndex: false,
       preInitialization: { testClient in
         testClient.handleMultipleRequests { (request: CreateWorkDoneProgressRequest) in
@@ -419,7 +417,7 @@ final class BackgroundIndexingTests: XCTestCase {
         """,
         "MyOtherFile.swift": "",
       ],
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
 
     let (uri, positions) = try project.openDocument("MyFile.swift")
@@ -486,7 +484,7 @@ final class BackgroundIndexingTests: XCTestCase {
         #include "Header.h"
         """,
       ],
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
 
     let (uri, positions) = try project.openDocument("Header.h", language: .c)
@@ -544,7 +542,7 @@ final class BackgroundIndexingTests: XCTestCase {
 
   func testPrepareTargetAfterEditToDependency() async throws {
     try await SkipUnless.swiftpmStoresModulesInSubdirectory()
-    var serverOptions = backgroundIndexingOptions
+    var serverOptions = SourceKitLSPServer.Options.testDefault
     let expectedPreparationTracker = ExpectedIndexTaskTracker(expectedPreparations: [
       [
         ExpectedPreparation(targetID: "LibA", runDestinationID: "dummy"),
@@ -580,6 +578,7 @@ final class BackgroundIndexingTests: XCTestCase {
         )
         """,
       serverOptions: serverOptions,
+      enableBackgroundIndexing: true,
       cleanUp: { expectedPreparationTracker.keepAlive() }
     )
 
@@ -637,7 +636,7 @@ final class BackgroundIndexingTests: XCTestCase {
     let libDPreparedForEditing = self.expectation(description: "LibD prepared for editing")
 
     try await SkipUnless.swiftpmStoresModulesInSubdirectory()
-    var serverOptions = backgroundIndexingOptions
+    var serverOptions = SourceKitLSPServer.Options.testDefault
     let expectedPreparationTracker = ExpectedIndexTaskTracker(expectedPreparations: [
       // Preparation of targets during the initial of the target
       [
@@ -689,6 +688,7 @@ final class BackgroundIndexingTests: XCTestCase {
         )
         """,
       serverOptions: serverOptions,
+      enableBackgroundIndexing: true,
       cleanUp: { expectedPreparationTracker.keepAlive() }
     )
 
@@ -718,7 +718,7 @@ final class BackgroundIndexingTests: XCTestCase {
       files: [
         "MyFile.swift": ""
       ],
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
     let targetPrepareNotification = try await project.testClient.nextNotification(ofType: LogMessageNotification.self)
     XCTAssert(
@@ -732,13 +732,13 @@ final class BackgroundIndexingTests: XCTestCase {
     )
   }
 
-  func testPreparationHappensInParallel() async throws {
+  func testIndexingHappensInParallel() async throws {
     try await SkipUnless.swiftpmStoresModulesInSubdirectory()
 
     let fileAIndexingStarted = self.expectation(description: "FileA indexing started")
     let fileBIndexingStarted = self.expectation(description: "FileB indexing started")
 
-    var serverOptions = backgroundIndexingOptions
+    var serverOptions = SourceKitLSPServer.Options.testDefault
     let expectedIndexTaskTracker = ExpectedIndexTaskTracker(
       expectedIndexStoreUpdates: [
         [
@@ -771,6 +771,7 @@ final class BackgroundIndexingTests: XCTestCase {
         "FileB.swift": "",
       ],
       serverOptions: serverOptions,
+      enableBackgroundIndexing: true,
       cleanUp: { expectedIndexTaskTracker.keepAlive() }
     )
   }
@@ -801,10 +802,10 @@ final class BackgroundIndexingTests: XCTestCase {
           ]
         )
         """,
-      serverOptions: backgroundIndexingOptions
+      enableBackgroundIndexing: true
     )
 
-    var otherClientOptions = backgroundIndexingOptions
+    var otherClientOptions = SourceKitLSPServer.Options.testDefault
     otherClientOptions.indexTestHooks = IndexTestHooks(
       preparationTaskDidStart: { taskDescription in
         XCTFail("Did not expect any target preparation, got \(taskDescription.targetsToPrepare)")
@@ -815,6 +816,7 @@ final class BackgroundIndexingTests: XCTestCase {
     )
     let otherClient = try await TestSourceKitLSPClient(
       serverOptions: otherClientOptions,
+      enableBackgroundIndexing: true,
       workspaceFolders: [
         WorkspaceFolder(uri: DocumentURI(project.scratchDirectory))
       ]