Show work done progress while a source file is being prepared for editor functionality

While `SemanticIndexManager.inProgressPrepareForEditorTask` is not `nil`, show a work done progress in the editor that the current file is being prepared for editor functionality.

I decided to use the indexing progress indicator for this for now because I think it’s too spammy if SourceKit-LSP has two different progress indicators for preparation and indexing. I’ll need to see how this feels like in practice.

rdar://128722609

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -64,6 +64,38 @@ public enum IndexTaskStatus: Comparable {
   case executing
 }
 
+/// The current index status that should be displayed to the editor.
+///
+/// In reality, these status are not exclusive. Eg. the index might be preparing one target for editor functionality,
+/// re-generating the build graph and indexing files at the same time. To avoid showing too many concurrent status
+/// messages to the user, we only show the highest priority task.
+public enum IndexProgressStatus {
+  case preparingFileForEditorFunctionality
+  case generatingBuildGraph
+  case indexing(preparationTasks: [ConfiguredTarget: IndexTaskStatus], indexTasks: [DocumentURI: IndexTaskStatus])
+  case upToDate
+
+  public func merging(with other: IndexProgressStatus) -> IndexProgressStatus {
+    switch (self, other) {
+    case (_, .preparingFileForEditorFunctionality), (.preparingFileForEditorFunctionality, _):
+      return .preparingFileForEditorFunctionality
+    case (_, .generatingBuildGraph), (.generatingBuildGraph, _):
+      return .generatingBuildGraph
+    case (
+      .indexing(let selfPreparationTasks, let selfIndexTasks),
+      .indexing(let otherPreparationTasks, let otherIndexTasks)
+    ):
+      return .indexing(
+        preparationTasks: selfPreparationTasks.merging(otherPreparationTasks) { max($0, $1) },
+        indexTasks: selfIndexTasks.merging(otherIndexTasks) { max($0, $1) }
+      )
+    case (.indexing, .upToDate): return self
+    case (.upToDate, .indexing): return other
+    case (.upToDate, .upToDate): return .upToDate
+    }
+  }
+}
+
 /// Schedules index tasks and keeps track of the index status of files.
 public final actor SemanticIndexManager {
   /// The underlying index. This is used to check if the index of a file is already up-to-date, in which case it doesn't
@@ -122,24 +154,22 @@ public final actor SemanticIndexManager {
   /// The parameter is the number of files that were scheduled to be indexed.
   private let indexTasksWereScheduled: @Sendable (_ numberOfFileScheduled: Int) -> Void
 
-  /// Callback that is called when the progress status of an update indexstore or preparation task finishes.
-  ///
-  /// An object observing this property probably wants to check `inProgressIndexTasks` when the callback is called to
-  /// get the current list of in-progress index tasks.
-  ///
-  /// The number of `indexStatusDidChange` calls does not have to relate to the number of `indexTasksWereScheduled` calls.
-  private let indexStatusDidChange: @Sendable () -> Void
+  /// Callback that is called when `progressStatus` might have changed.
+  private let indexProgressStatusDidChange: @Sendable () -> Void
 
   // MARK: - Public API
 
   /// A summary of the tasks that this `SemanticIndexManager` has currently scheduled or is currently indexing.
-  public var inProgressTasks:
-    (
-      isGeneratingBuildGraph: Bool,
-      indexTasks: [DocumentURI: IndexTaskStatus],
-      preparationTasks: [ConfiguredTarget: IndexTaskStatus]
-    )
-  {
+  public var progressStatus: IndexProgressStatus {
+    if inProgressPrepareForEditorTask != nil {
+      return .preparingFileForEditorFunctionality
+    }
+    if generateBuildGraphTask != nil {
+      return .generatingBuildGraph
+    }
+    let preparationTasks = inProgressPreparationTasks.mapValues { queuedTask in
+      return queuedTask.isExecuting ? IndexTaskStatus.executing : IndexTaskStatus.scheduled
+    }
     let indexTasks = inProgressIndexTasks.mapValues { status in
       switch status {
       case .waitingForPreparation:
@@ -148,10 +178,10 @@ public final actor SemanticIndexManager {
         return updateIndexStoreTask.isExecuting ? IndexTaskStatus.executing : IndexTaskStatus.scheduled
       }
     }
-    let preparationTasks = inProgressPreparationTasks.mapValues { queuedTask in
-      return queuedTask.isExecuting ? IndexTaskStatus.executing : IndexTaskStatus.scheduled
+    if preparationTasks.isEmpty && indexTasks.isEmpty {
+      return .upToDate
     }
-    return (generateBuildGraphTask != nil, indexTasks, preparationTasks)
+    return .indexing(preparationTasks: preparationTasks, indexTasks: indexTasks)
   }
 
   public init(
@@ -161,7 +191,7 @@ public final actor SemanticIndexManager {
     indexTaskScheduler: TaskScheduler<AnyIndexTaskDescription>,
     indexProcessDidProduceResult: @escaping @Sendable (IndexProcessResult) -> Void,
     indexTasksWereScheduled: @escaping @Sendable (Int) -> Void,
-    indexStatusDidChange: @escaping @Sendable () -> Void
+    indexProgressStatusDidChange: @escaping @Sendable () -> Void
   ) {
     self.index = index
     self.buildSystemManager = buildSystemManager
@@ -169,7 +199,7 @@ public final actor SemanticIndexManager {
     self.indexTaskScheduler = indexTaskScheduler
     self.indexProcessDidProduceResult = indexProcessDidProduceResult
     self.indexTasksWereScheduled = indexTasksWereScheduled
-    self.indexStatusDidChange = indexStatusDidChange
+    self.indexProgressStatusDidChange = indexProgressStatusDidChange
   }
 
   /// Schedules a task to index `files`. Files that are known to be up-to-date based on `indexStatus` will
@@ -222,7 +252,7 @@ public final actor SemanticIndexManager {
         generateBuildGraphTask = nil
       }
     }
-    indexStatusDidChange()
+    indexProgressStatusDidChange()
   }
 
   /// Wait for all in-progress index tasks to finish.
@@ -350,11 +380,13 @@ public final actor SemanticIndexManager {
         await self.prepare(targets: [target], priority: priority)
         if inProgressPrepareForEditorTask?.id == id {
           inProgressPrepareForEditorTask = nil
+          self.indexProgressStatusDidChange()
         }
       }
     }
     inProgressPrepareForEditorTask?.task.cancel()
     inProgressPrepareForEditorTask = (id, uri, task)
+    self.indexProgressStatusDidChange()
   }
 
   // MARK: - Helper functions
@@ -388,7 +420,7 @@ public final actor SemanticIndexManager {
     }
     let preparationTask = await indexTaskScheduler.schedule(priority: priority, taskDescription) { task, newState in
       guard case .finished = newState else {
-        self.indexStatusDidChange()
+        self.indexProgressStatusDidChange()
         return
       }
       for target in targetsToPrepare {
@@ -396,7 +428,7 @@ public final actor SemanticIndexManager {
           self.inProgressPreparationTasks[target] = nil
         }
       }
-      self.indexStatusDidChange()
+      self.indexProgressStatusDidChange()
     }
     for target in targetsToPrepare {
       inProgressPreparationTasks[target] = OpaqueQueuedIndexTask(preparationTask)
@@ -432,7 +464,7 @@ public final actor SemanticIndexManager {
     )
     let updateIndexTask = await indexTaskScheduler.schedule(priority: priority, taskDescription) { task, newState in
       guard case .finished = newState else {
-        self.indexStatusDidChange()
+        self.indexProgressStatusDidChange()
         return
       }
       for fileAndTarget in filesAndTargets {
@@ -442,7 +474,7 @@ public final actor SemanticIndexManager {
           self.inProgressIndexTasks[fileAndTarget.file.sourceFile] = nil
         }
       }
-      self.indexStatusDidChange()
+      self.indexProgressStatusDidChange()
     }
     for fileAndTarget in filesAndTargets {
       if case .waitingForPreparation(preparationTaskID, let indexTask) = inProgressIndexTasks[

Sources/SourceKitLSP/IndexProgressManager.swift

@@ -18,11 +18,11 @@ import SemanticIndex
 /// Listens for index status updates from `SemanticIndexManagers`. From that information, it manages a
 /// `WorkDoneProgress` that communicates the index progress to the editor.
 actor IndexProgressManager {
-  /// A queue on which `indexTaskWasQueued` and `indexStatusDidChange` are handled.
+  /// A queue on which `indexTaskWasQueued` and `indexProgressStatusDidChange` are handled.
   ///
-  /// This allows the two functions two be `nonisolated` (and eg. the caller of `indexStatusDidChange` doesn't have to
+  /// This allows the two functions two be `nonisolated` (and eg. the caller of `indexProgressStatusDidChange` doesn't have to
   /// wait for the work done progress to be updated) while still guaranteeing that there is only one
-  /// `indexStatusDidChangeImpl` running at a time, preventing race conditions that would cause two
+  /// `indexProgressStatusDidChangeImpl` running at a time, preventing race conditions that would cause two
   /// `WorkDoneProgressManager`s to be created.
   private let queue = AsyncQueue<Serial>()
 
@@ -64,74 +64,64 @@ actor IndexProgressManager {
 
   private func indexTasksWereScheduledImpl(count: Int) async {
     queuedIndexTasks += count
-    await indexStatusDidChangeImpl()
+    await indexProgressStatusDidChangeImpl()
   }
 
   /// Called when a `SemanticIndexManager` finishes indexing a file. Adjusts the done index count, eg. the 1 in `1/3`.
-  nonisolated func indexStatusDidChange() {
+  nonisolated func indexProgressStatusDidChange() {
     queue.async {
-      await self.indexStatusDidChangeImpl()
+      await self.indexProgressStatusDidChangeImpl()
     }
   }
 
-  private func indexStatusDidChangeImpl() async {
+  private func indexProgressStatusDidChangeImpl() async {
     guard let sourceKitLSPServer else {
       workDoneProgress = nil
       return
     }
-    var isGeneratingBuildGraph = false
-    var indexTasks: [DocumentURI: IndexTaskStatus] = [:]
-    var preparationTasks: [ConfiguredTarget: IndexTaskStatus] = [:]
+    var status = IndexProgressStatus.upToDate
     for indexManager in await sourceKitLSPServer.workspaces.compactMap({ $0.semanticIndexManager }) {
-      let inProgress = await indexManager.inProgressTasks
-      isGeneratingBuildGraph = isGeneratingBuildGraph || inProgress.isGeneratingBuildGraph
-      indexTasks.merge(inProgress.indexTasks) { lhs, rhs in
-        return max(lhs, rhs)
-      }
-      preparationTasks.merge(inProgress.preparationTasks) { lhs, rhs in
-        return max(lhs, rhs)
-      }
+      status = status.merging(with: await indexManager.progressStatus)
     }
 
-    if indexTasks.isEmpty && !isGeneratingBuildGraph {
+    var message: String
+    let percentage: Int
+    switch status {
+    case .preparingFileForEditorFunctionality:
+      message = "Preparing current file"
+      percentage = 0
+    case .generatingBuildGraph:
+      message = "Generating build graph"
+      percentage = 0
+    case .indexing(preparationTasks: let preparationTasks, indexTasks: let indexTasks):
+      // We can get into a situation where queuedIndexTasks < indexTasks.count if we haven't processed all
+      // `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)
+      message = "\(finishedTasks) / \(queuedIndexTasks)"
+      if await sourceKitLSPServer.options.indexOptions.showActivePreparationTasksInProgress {
+        var inProgressTasks: [String] = []
+        inProgressTasks += preparationTasks.filter { $0.value == .executing }
+          .map { "- Preparing \($0.key.targetID)" }
+          .sorted()
+        inProgressTasks += indexTasks.filter { $0.value == .executing }
+          .map { "- Indexing \($0.key.fileURL?.lastPathComponent ?? $0.key.pseudoPath)" }
+          .sorted()
+
+        message += "\n\n" + inProgressTasks.joined(separator: "\n")
+      }
+      if queuedIndexTasks != 0 {
+        percentage = Int(Double(finishedTasks) / Double(queuedIndexTasks) * 100)
+      } else {
+        percentage = 0
+      }
+    case .upToDate:
       // Nothing left to index. Reset the target count and dismiss the work done progress.
       queuedIndexTasks = 0
       workDoneProgress = nil
       return
     }
 
-    // We can get into a situation where queuedIndexTasks < indexTasks.count if we haven't processed all
-    // `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: String
-    if isGeneratingBuildGraph {
-      message = "Generating build graph"
-    } else {
-      message = "\(finishedTasks) / \(queuedIndexTasks)"
-    }
-
-    if await sourceKitLSPServer.options.indexOptions.showActivePreparationTasksInProgress {
-      var inProgressTasks: [String] = []
-      if isGeneratingBuildGraph {
-        inProgressTasks.append("- Generating build graph")
-      }
-      inProgressTasks += preparationTasks.filter { $0.value == .executing }
-        .map { "- Preparing \($0.key.targetID)" }
-        .sorted()
-      inProgressTasks += indexTasks.filter { $0.value == .executing }
-        .map { "- Indexing \($0.key.fileURL?.lastPathComponent ?? $0.key.pseudoPath)" }
-        .sorted()
-
-      message += "\n\n" + inProgressTasks.joined(separator: "\n")
-    }
-
-    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 {

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -510,7 +510,7 @@ public actor SourceKitLSPServer {
       uriToWorkspaceCache = [:]
       // `indexProgressManager` iterates over all workspaces in the SourceKitLSPServer. Modifying workspaces might thus
       // update the index progress status.
-      indexProgressManager.indexStatusDidChange()
+      indexProgressManager.indexProgressStatusDidChange()
     }
   }
 
@@ -1257,8 +1257,8 @@ extension SourceKitLSPServer {
       indexTasksWereScheduled: { [weak self] count in
         self?.indexProgressManager.indexTasksWereScheduled(count: count)
       },
-      indexStatusDidChange: { [weak self] in
-        self?.indexProgressManager.indexStatusDidChange()
+      indexProgressStatusDidChange: { [weak self] in
+        self?.indexProgressManager.indexProgressStatusDidChange()
       }
     )
   }
@@ -1323,8 +1323,8 @@ extension SourceKitLSPServer {
           indexTasksWereScheduled: { [weak self] count in
             self?.indexProgressManager.indexTasksWereScheduled(count: count)
           },
-          indexStatusDidChange: { [weak self] in
-            self?.indexProgressManager.indexStatusDidChange()
+          indexProgressStatusDidChange: { [weak self] in
+            self?.indexProgressManager.indexProgressStatusDidChange()
           }
         )
 

Sources/SourceKitLSP/WorkDoneProgressManager.swift

@@ -36,6 +36,10 @@ final class WorkDoneProgressManager {
   ///    - This should have `workDoneProgressCreated == true` so that it can send the work progress end.
   private let workDoneProgressCreated: ThreadSafeBox<Bool> & AnyObject = ThreadSafeBox<Bool>(initialValue: false)
 
+  /// The last message and percentage so we don't send a new report notification to the client if `update` is called
+  /// without any actual change.
+  private var lastStatus: (message: String?, percentage: Int?)
+
   convenience init?(server: SourceKitLSPServer, title: String, message: String? = nil, percentage: Int? = nil) async {
     guard let capabilityRegistry = await server.capabilityRegistry else {
       return nil
@@ -69,6 +73,7 @@ final class WorkDoneProgressManager {
         )
       )
       workDoneProgressCreated.value = true
+      self.lastStatus = (message, percentage)
     }
   }
 
@@ -77,6 +82,10 @@ final class WorkDoneProgressManager {
       guard workDoneProgressCreated.value else {
         return
       }
+      guard (message, percentage) != self.lastStatus else {
+        return
+      }
+      self.lastStatus = (message, percentage)
       server.sendNotificationToClient(
         WorkDoneProgress(
           token: token,

Sources/SourceKitLSP/Workspace.swift

@@ -96,7 +96,7 @@ public final class Workspace: Sendable {
     indexTaskScheduler: TaskScheduler<AnyIndexTaskDescription>,
     indexProcessDidProduceResult: @escaping @Sendable (IndexProcessResult) -> Void,
     indexTasksWereScheduled: @escaping @Sendable (Int) -> Void,
-    indexStatusDidChange: @escaping @Sendable () -> Void
+    indexProgressStatusDidChange: @escaping @Sendable () -> Void
   ) async {
     self.documentManager = documentManager
     self.buildSetup = options.buildSetup
@@ -117,7 +117,7 @@ public final class Workspace: Sendable {
         indexTaskScheduler: indexTaskScheduler,
         indexProcessDidProduceResult: indexProcessDidProduceResult,
         indexTasksWereScheduled: indexTasksWereScheduled,
-        indexStatusDidChange: indexStatusDidChange
+        indexProgressStatusDidChange: indexProgressStatusDidChange
       )
     } else {
       self.semanticIndexManager = nil
@@ -156,7 +156,7 @@ public final class Workspace: Sendable {
     indexProcessDidProduceResult: @escaping @Sendable (IndexProcessResult) -> Void,
     reloadPackageStatusCallback: @Sendable @escaping (ReloadPackageStatus) async -> Void,
     indexTasksWereScheduled: @Sendable @escaping (Int) -> Void,
-    indexStatusDidChange: @Sendable @escaping () -> Void
+    indexProgressStatusDidChange: @Sendable @escaping () -> Void
   ) async throws {
     var buildSystem: BuildSystem? = nil
 
@@ -263,7 +263,7 @@ public final class Workspace: Sendable {
       indexTaskScheduler: indexTaskScheduler,
       indexProcessDidProduceResult: indexProcessDidProduceResult,
       indexTasksWereScheduled: indexTasksWereScheduled,
-      indexStatusDidChange: indexStatusDidChange
+      indexProgressStatusDidChange: indexProgressStatusDidChange
     )
   }
 

Tests/SourceKitLSPTests/BackgroundIndexingTests.swift

@@ -579,6 +579,7 @@ final class BackgroundIndexingTests: XCTestCase {
           ]
         )
         """,
+      capabilities: ClientCapabilities(window: WindowClientCapabilities(workDoneProgress: true)),
       serverOptions: serverOptions,
       cleanUp: { expectedPreparationTracker.keepAlive() }
     )
@@ -604,6 +605,9 @@ final class BackgroundIndexingTests: XCTestCase {
     )
 
     let receivedEmptyDiagnostics = self.expectation(description: "Received diagnostic refresh request")
+    project.testClient.handleMultipleRequests { (_: CreateWorkDoneProgressRequest) in
+      return VoidResponse()
+    }
 
     project.testClient.handleMultipleRequests { (_: DiagnosticsRefreshRequest) in
       Task {
@@ -629,6 +633,18 @@ final class BackgroundIndexingTests: XCTestCase {
     )
 
     try await fulfillmentOfOrThrow([receivedEmptyDiagnostics])
+
+    // Check that we received a work done progress for the re-preparation of the target
+    _ = try await project.testClient.nextNotification(
+      ofType: WorkDoneProgress.self,
+      satisfying: { notification in
+        switch notification.value {
+        case .begin(let value): return value.message == "Preparing current file"
+        case .report(let value): return value.message == "Preparing current file"
+        case .end: return false
+        }
+      }
+    )
   }
 
   func testDontStackTargetPreparationForEditorFunctionality() async throws {

Tests/SourceKitLSPTests/BuildSystemTests.swift

@@ -141,7 +141,7 @@ final class BuildSystemTests: XCTestCase {
       indexTaskScheduler: .forTesting,
       indexProcessDidProduceResult: { _ in },
       indexTasksWereScheduled: { _ in },
-      indexStatusDidChange: {}
+      indexProgressStatusDidChange: {}
     )
 
     await server.setWorkspaces([(workspace: workspace, isImplicit: false)])