Show a work done progress while the semantic index is rebuilding a build graph

Sources/SKTestSupport/TestSourceKitLSPClient.swift

@@ -248,15 +248,16 @@ public final class TestSourceKitLSPClient: MessageHandler {
     return try await nextNotification(ofType: PublishDiagnosticsNotification.self, timeout: timeout)
   }
 
-  /// Waits for the next notification of the given type to be sent to the client. Ignores any notifications that are of
-  /// a different type.
+  /// Waits for the next notification of the given type to be sent to the client that satisfies the given predicate.
+  /// Ignores any notifications that are of a different type or that don't satisfy the predicate.
   public func nextNotification<ExpectedNotificationType: NotificationType>(
     ofType: ExpectedNotificationType.Type,
+    satisfying predicate: (ExpectedNotificationType) -> Bool = { _ in true },
     timeout: TimeInterval = defaultTimeout
   ) async throws -> ExpectedNotificationType {
     while true {
       let nextNotification = try await nextNotification(timeout: timeout)
-      if let notification = nextNotification as? ExpectedNotificationType {
+      if let notification = nextNotification as? ExpectedNotificationType, predicate(notification) {
         return notification
       }
     }

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -192,7 +192,9 @@ public final actor SemanticIndexManager {
       defer {
         signposter.endInterval("Preparing", state)
       }
+      await testHooks.buildGraphGenerationDidStart?()
       await orLog("Generating build graph") { try await self.buildSystemManager.generateBuildGraph() }
+      await testHooks.buildGraphGenerationDidFinish?()
       let index = index.checked(for: .modifiedFiles)
       let filesToIndex = await self.buildSystemManager.sourceFiles().lazy.map(\.uri)
         .filter { uri in
@@ -205,6 +207,7 @@ public final actor SemanticIndexManager {
       await scheduleBackgroundIndex(files: filesToIndex)
       generateBuildGraphTask = nil
     }
+    indexStatusDidChange()
   }
 
   /// Wait for all in-progress index tasks to finish.

Sources/SemanticIndex/TestHooks.swift

@@ -12,6 +12,10 @@
 
 /// Callbacks that allow inspection of internal state modifications during testing.
 public struct IndexTestHooks: Sendable {
+  public var buildGraphGenerationDidStart: (@Sendable () async -> Void)?
+
+  public var buildGraphGenerationDidFinish: (@Sendable () async -> Void)?
+
   public var preparationTaskDidStart: (@Sendable (PreparationTaskDescription) async -> Void)?
 
   public var preparationTaskDidFinish: (@Sendable (PreparationTaskDescription) async -> Void)?
@@ -22,11 +26,15 @@ public struct IndexTestHooks: Sendable {
   public var updateIndexStoreTaskDidFinish: (@Sendable (UpdateIndexStoreTaskDescription) async -> Void)?
 
   public init(
+    buildGraphGenerationDidStart: (@Sendable () async -> Void)? = nil,
+    buildGraphGenerationDidFinish: (@Sendable () async -> Void)? = nil,
     preparationTaskDidStart: (@Sendable (PreparationTaskDescription) async -> Void)? = nil,
     preparationTaskDidFinish: (@Sendable (PreparationTaskDescription) async -> Void)? = nil,
     updateIndexStoreTaskDidStart: (@Sendable (UpdateIndexStoreTaskDescription) async -> Void)? = nil,
     updateIndexStoreTaskDidFinish: (@Sendable (UpdateIndexStoreTaskDescription) async -> Void)? = nil
   ) {
+    self.buildGraphGenerationDidStart = buildGraphGenerationDidStart
+    self.buildGraphGenerationDidFinish = buildGraphGenerationDidFinish
     self.preparationTaskDidStart = preparationTaskDidStart
     self.preparationTaskDidFinish = preparationTaskDidFinish
     self.updateIndexStoreTaskDidStart = updateIndexStoreTaskDidStart

Sources/SourceKitLSP/IndexProgressManager.swift

@@ -93,7 +93,7 @@ actor IndexProgressManager {
       }
     }
 
-    if indexTasks.isEmpty {
+    if indexTasks.isEmpty && !isGeneratingBuildGraph {
       // Nothing left to index. Reset the target count and dismiss the work done progress.
       queuedIndexTasks = 0
       workDoneProgress = nil
@@ -104,7 +104,12 @@ actor IndexProgressManager {
     // `indexTasksWereScheduled` calls yet but the semantic index managers already track them in their in-progress tasks.
     // Clip the finished tasks to 0 because showing a negative number there looks stupid.
     let finishedTasks = max(queuedIndexTasks - indexTasks.count, 0)
-    var message = "\(finishedTasks) / \(queuedIndexTasks)"
+    var message: String
+    if isGeneratingBuildGraph {
+      message = "Generating build graph"
+    } else {
+      message = "\(finishedTasks) / \(queuedIndexTasks)"
+    }
 
     if await sourceKitLSPServer.options.indexOptions.showActivePreparationTasksInProgress {
       var inProgressTasks: [String] = []
@@ -121,7 +126,12 @@ actor IndexProgressManager {
       message += "\n\n" + inProgressTasks.joined(separator: "\n")
     }
 
-    let percentage = Int(Double(finishedTasks) / Double(queuedIndexTasks) * 100)
+    let percentage: Int
+    if queuedIndexTasks != 0 {
+      percentage = Int(Double(finishedTasks) / Double(queuedIndexTasks) * 100)
+    } else {
+      percentage = 0
+    }
     if let workDoneProgress {
       workDoneProgress.update(message: message, percentage: percentage)
     } else {

Tests/SourceKitLSPTests/BackgroundIndexingTests.swift

@@ -333,6 +333,21 @@ final class BackgroundIndexingTests: XCTestCase {
   }
 
   func testBackgroundIndexingStatusWorkDoneProgress() async throws {
+    let receivedBeginProgressNotification = self.expectation(
+      description: "Received work done progress saying build graph generation"
+    )
+    let receivedReportProgressNotification = self.expectation(
+      description: "Received work done progress saying indexing"
+    )
+    var serverOptions = backgroundIndexingOptions
+    serverOptions.indexTestHooks = IndexTestHooks(
+      buildGraphGenerationDidFinish: {
+        await self.fulfillment(of: [receivedBeginProgressNotification], timeout: defaultTimeout)
+      },
+      updateIndexStoreTaskDidFinish: { _ in
+        await self.fulfillment(of: [receivedReportProgressNotification], timeout: defaultTimeout)
+      }
+    )
     let project = try await SwiftPMTestProject(
       files: [
         "MyFile.swift": """
@@ -343,36 +358,57 @@ final class BackgroundIndexingTests: XCTestCase {
         """
       ],
       capabilities: ClientCapabilities(window: WindowClientCapabilities(workDoneProgress: true)),
-      serverOptions: backgroundIndexingOptions,
+      serverOptions: serverOptions,
+      pollIndex: false,
       preInitialization: { testClient in
         testClient.handleMultipleRequests { (request: CreateWorkDoneProgressRequest) in
           return VoidResponse()
         }
       }
     )
-    var indexingWorkDoneProgressToken: ProgressToken? = nil
-    var didGetEndWorkDoneProgress = false
-    // Loop terminates when we see the work done end progress or if waiting for the next notification times out
-    LOOP: while true {
-      let workDoneProgress = try await project.testClient.nextNotification(ofType: WorkDoneProgress.self)
-      switch workDoneProgress.value {
-      case .begin(let data):
-        if data.title == "Indexing" {
-          XCTAssertNil(indexingWorkDoneProgressToken, "Received multiple work done progress notifications for indexing")
-          indexingWorkDoneProgressToken = workDoneProgress.token
+
+    let beginNotification = try await project.testClient.nextNotification(
+      ofType: WorkDoneProgress.self,
+      satisfying: { notification in
+        guard case .begin(let data) = notification.value else {
+          return false
         }
-      case .report:
-        // We ignore progress reports in the test because it's non-deterministic how many we get
-        break
-      case .end:
-        if workDoneProgress.token == indexingWorkDoneProgressToken {
-          didGetEndWorkDoneProgress = true
-          break LOOP
+        return data.title == "Indexing"
+      }
+    )
+    receivedBeginProgressNotification.fulfill()
+    guard case .begin(let beginData) = beginNotification.value else {
+      XCTFail("Expected begin notification")
+      return
+    }
+    XCTAssertEqual(beginData.message, "Generating build graph")
+    let indexingWorkDoneProgressToken = beginNotification.token
+
+    let reportNotification = try await project.testClient.nextNotification(
+      ofType: WorkDoneProgress.self,
+      satisfying: { notification in
+        guard notification.token == indexingWorkDoneProgressToken, case .report = notification.value else {
+          return false
         }
+        return true
       }
+    )
+    receivedReportProgressNotification.fulfill()
+    guard case .report(let reportData) = reportNotification.value else {
+      XCTFail("Expected report notification")
+      return
     }
-    XCTAssertNotNil(indexingWorkDoneProgressToken, "Expected to receive a work done progress start")
-    XCTAssert(didGetEndWorkDoneProgress, "Expected end work done progress")
+    XCTAssertEqual(reportData.message, "0 / 1")
+
+    _ = try await project.testClient.nextNotification(
+      ofType: WorkDoneProgress.self,
+      satisfying: { notification in
+        guard notification.token == indexingWorkDoneProgressToken, case .end = notification.value else {
+          return false
+        }
+        return true
+      }
+    )
 
     withExtendedLifetime(project) {}
   }