Add a work done progress that shows the index progress

Sources/SKCore/TaskScheduler.swift

@@ -76,6 +76,19 @@ public protocol TaskDescriptionProtocol: Identifiable, Sendable, CustomLogString
   var estimatedCPUCoreCount: Int { get }
 }
 
+/// Parameter that's passed to `executionStateChangedCallback` to indicate the new state of a scheduled task.
+public enum TaskExecutionState {
+  /// The task started executing.
+  case executing
+
+  /// The task was cancelled and will be re-scheduled for execution later. Will be followed by another call with
+  /// `executing`.
+  case cancelledToBeRescheduled
+
+  /// The task has finished executing. Now more state updates will come after this one.
+  case finished
+}
+
 fileprivate actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
   /// Result of `executionTask` / the tasks in `executionTaskCreatedContinuation`.
   /// See doc comment on `executionTask`.
@@ -136,9 +149,18 @@ fileprivate actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
   /// Gets reset every time `executionTask` finishes.
   nonisolated(unsafe) private var cancelledToBeRescheduled: AtomicBool = .init(initialValue: false)
 
-  init(priority: TaskPriority? = nil, description: TaskDescription) async {
+  /// A callback that will be called when the task starts executing, is cancelled to be rescheduled, or when it finishes
+  /// execution.
+  private let executionStateChangedCallback: (@Sendable (TaskExecutionState) async -> Void)?
+
+  init(
+    priority: TaskPriority? = nil,
+    description: TaskDescription,
+    executionStateChangedCallback: (@Sendable (TaskExecutionState) async -> Void)?
+  ) async {
     self._priority = .init(initialValue: priority?.rawValue ?? Task.currentPriority.rawValue)
     self.description = description
+    self.executionStateChangedCallback = executionStateChangedCallback
 
     var updatePriorityContinuation: AsyncStream<Void>.Continuation!
     let updatePriorityStream = AsyncStream {
@@ -194,16 +216,19 @@ fileprivate actor QueuedTask<TaskDescription: TaskDescriptionProtocol> {
     }
     executionTask = task
     executionTaskCreatedContinuation.yield(task)
+    await executionStateChangedCallback?(.executing)
     return await task.value
   }
 
   /// Implementation detail of `execute` that is called after `self.description.execute()` finishes.
-  private func finalizeExecution() -> ExecutionTaskFinishStatus {
+  private func finalizeExecution() async -> ExecutionTaskFinishStatus {
     self.executionTask = nil
     if Task.isCancelled && self.cancelledToBeRescheduled.value {
+      await executionStateChangedCallback?(.cancelledToBeRescheduled)
       self.cancelledToBeRescheduled.value = false
       return ExecutionTaskFinishStatus.cancelledToBeRescheduled
     } else {
+      await executionStateChangedCallback?(.finished)
       return ExecutionTaskFinishStatus.terminated
     }
   }
@@ -308,12 +333,17 @@ public actor TaskScheduler<TaskDescription: TaskDescriptionProtocol> {
   @discardableResult
   public func schedule(
     priority: TaskPriority? = nil,
-    _ taskDescription: TaskDescription
+    _ taskDescription: TaskDescription,
+    @_inheritActorContext executionStateChangedCallback: (@Sendable (TaskExecutionState) async -> Void)? = nil
   ) async -> Task<Void, Never> {
     withLoggingSubsystemAndScope(subsystem: taskSchedulerSubsystem, scope: nil) {
       logger.debug("Scheduling \(taskDescription.forLogging)")
     }
-    let queuedTask = await QueuedTask(priority: priority, description: taskDescription)
+    let queuedTask = await QueuedTask(
+      priority: priority,
+      description: taskDescription,
+      executionStateChangedCallback: executionStateChangedCallback
+    )
     pendingTasks.append(queuedTask)
     Task.detached(priority: priority ?? Task.currentPriority) {
       // Poke the `TaskScheduler` to execute a new task. If the `TaskScheduler` is already working at its capacity

Sources/SKTestSupport/MultiFileTestProject.swift

@@ -80,8 +80,10 @@ public class MultiFileTestProject {
   public init(
     files: [RelativeFileLocation: String],
     workspaces: (URL) async throws -> [WorkspaceFolder] = { [WorkspaceFolder(uri: DocumentURI($0))] },
+    capabilities: ClientCapabilities = ClientCapabilities(),
     serverOptions: SourceKitLSPServer.Options = .testDefault,
     usePullDiagnostics: Bool = true,
+    preInitialization: ((TestSourceKitLSPClient) -> Void)? = nil,
     testName: String = #function
   ) async throws {
     scratchDirectory = try testScratchDir(testName: testName)
@@ -112,8 +114,10 @@ public class MultiFileTestProject {
 
     self.testClient = try await TestSourceKitLSPClient(
       serverOptions: serverOptions,
+      capabilities: capabilities,
       usePullDiagnostics: usePullDiagnostics,
       workspaceFolders: workspaces(scratchDirectory),
+      preInitialization: preInitialization,
       cleanUp: { [scratchDirectory] in
         if cleanScratchDirectories {
           try? FileManager.default.removeItem(at: scratchDirectory)

Sources/SKTestSupport/SwiftPMTestProject.swift

@@ -42,8 +42,10 @@ public class SwiftPMTestProject: MultiFileTestProject {
     workspaces: (URL) async throws -> [WorkspaceFolder] = { [WorkspaceFolder(uri: DocumentURI($0))] },
     build: Bool = false,
     allowBuildFailure: Bool = false,
+    capabilities: ClientCapabilities = ClientCapabilities(),
     serverOptions: SourceKitLSPServer.Options = .testDefault,
     pollIndex: Bool = true,
+    preInitialization: ((TestSourceKitLSPClient) -> Void)? = nil,
     usePullDiagnostics: Bool = true,
     testName: String = #function
   ) async throws {
@@ -66,8 +68,10 @@ public class SwiftPMTestProject: MultiFileTestProject {
     try await super.init(
       files: filesByPath,
       workspaces: workspaces,
+      capabilities: capabilities,
       serverOptions: serverOptions,
       usePullDiagnostics: usePullDiagnostics,
+      preInitialization: preInitialization,
       testName: testName
     )
 

Sources/SKTestSupport/TestSourceKitLSPClient.swift

@@ -83,6 +83,8 @@ public final class TestSourceKitLSPClient: MessageHandler {
   ///   - capabilities: The test client's capabilities.
   ///   - usePullDiagnostics: Whether to use push diagnostics or use push-based diagnostics
   ///   - 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.
   ///   - cleanUp: A closure that is called when the `TestSourceKitLSPClient` is destructed.
   ///     This allows e.g. a `IndexedSingleSwiftFileTestProject` to delete its temporary files when they are no longer
   ///     needed.
@@ -94,6 +96,7 @@ public final class TestSourceKitLSPClient: MessageHandler {
     capabilities: ClientCapabilities = ClientCapabilities(),
     usePullDiagnostics: Bool = true,
     workspaceFolders: [WorkspaceFolder]? = nil,
+    preInitialization: ((TestSourceKitLSPClient) -> Void)? = nil,
     cleanUp: @escaping () -> Void = {}
   ) async throws {
     if !useGlobalModuleCache {
@@ -135,7 +138,7 @@ public final class TestSourceKitLSPClient: MessageHandler {
       guard capabilities.textDocument!.diagnostic == nil else {
         struct ConflictingDiagnosticsError: Error, CustomStringConvertible {
           var description: String {
-            "usePushDiagnostics = false is not supported if capabilities already contain diagnostic options"
+            "usePullDiagnostics = false is not supported if capabilities already contain diagnostic options"
           }
         }
         throw ConflictingDiagnosticsError()
@@ -145,6 +148,7 @@ public final class TestSourceKitLSPClient: MessageHandler {
         XCTAssertEqual(request.registrations.only?.method, DocumentDiagnosticsRequest.method)
         return VoidResponse()
       }
+      preInitialization?(self)
     }
     if initialize {
       _ = try await self.send(
@@ -286,18 +290,19 @@ public final class TestSourceKitLSPClient: MessageHandler {
     id: LanguageServerProtocol.RequestID,
     reply: @escaping (LSPResult<Request.Response>) -> Void
   ) {
-    guard let requestHandler = requestHandlers.first else {
-      reply(.failure(.methodNotFound(Request.method)))
-      return
-    }
-    guard let requestHandler = requestHandler as? RequestHandler<Request> else {
-      print("\(RequestHandler<Request>.self)")
-      XCTFail("Received request of unexpected type \(Request.method)")
+    let requestHandlerAndIndex = requestHandlers.enumerated().compactMap {
+      (index, handler) -> (RequestHandler<Request>, Int)? in
+      guard let handler = handler as? RequestHandler<Request> else {
+        return nil
+      }
+      return (handler, index)
+    }.first
+    guard let (requestHandler, index) = requestHandlerAndIndex else {
       reply(.failure(.methodNotFound(Request.method)))
       return
     }
     reply(.success(requestHandler(params)))
-    requestHandlers.removeFirst()
+    requestHandlers.remove(at: index)
   }
 
   // MARK: - Convenience functions

Sources/SemanticIndex/IndexTaskDescription.swift

@@ -12,94 +12,72 @@
 
 import SKCore
 
-/// A task that either prepares targets or updates the index store for a set of files.
-public enum IndexTaskDescription: TaskDescriptionProtocol {
-  case updateIndexStore(UpdateIndexStoreTaskDescription)
-  case preparation(PreparationTaskDescription)
+/// Protocol of tasks that are executed on the index task scheduler.
+///
+/// It is assumed that `IndexTaskDescription` of different types are allowed to execute in parallel.
+protocol IndexTaskDescription: TaskDescriptionProtocol {
+  /// A string that is unique to this type of `IndexTaskDescription`. It is used to produce unique IDs for tasks of
+  /// different types in `AnyIndexTaskDescription`
+  static var idPrefix: String { get }
+
+  var id: UInt32 { get }
+}
+
+extension IndexTaskDescription {
+  func dependencies(
+    to currentlyExecutingTasks: [AnyIndexTaskDescription]
+  ) -> [TaskDependencyAction<AnyIndexTaskDescription>] {
+    return self.dependencies(to: currentlyExecutingTasks.compactMap { $0.wrapped as? Self })
+      .map {
+        switch $0 {
+        case .cancelAndRescheduleDependency(let td):
+          return .cancelAndRescheduleDependency(AnyIndexTaskDescription(td))
+        case .waitAndElevatePriorityOfDependency(let td):
+          return .waitAndElevatePriorityOfDependency(AnyIndexTaskDescription(td))
+        }
+      }
+
+  }
+}
+
+/// Type-erased wrapper of an `IndexTaskDescription`.
+public struct AnyIndexTaskDescription: TaskDescriptionProtocol {
+  let wrapped: any IndexTaskDescription
+
+  init(_ wrapped: any IndexTaskDescription) {
+    self.wrapped = wrapped
+  }
 
   public var isIdempotent: Bool {
-    switch self {
-    case .updateIndexStore(let taskDescription): return taskDescription.isIdempotent
-    case .preparation(let taskDescription): return taskDescription.isIdempotent
-    }
+    return wrapped.isIdempotent
   }
 
   public var estimatedCPUCoreCount: Int {
-    switch self {
-    case .updateIndexStore(let taskDescription): return taskDescription.estimatedCPUCoreCount
-    case .preparation(let taskDescription): return taskDescription.estimatedCPUCoreCount
-    }
+    return wrapped.estimatedCPUCoreCount
   }
 
   public var id: String {
-    switch self {
-    case .updateIndexStore(let taskDescription): return "indexing-\(taskDescription.id)"
-    case .preparation(let taskDescription): return "preparation-\(taskDescription.id)"
-    }
+    return "\(type(of: wrapped).idPrefix)-\(wrapped.id)"
   }
 
   public var description: String {
-    switch self {
-    case .updateIndexStore(let taskDescription): return taskDescription.description
-    case .preparation(let taskDescription): return taskDescription.description
-    }
+    return wrapped.description
   }
 
   public var redactedDescription: String {
-    switch self {
-    case .updateIndexStore(let taskDescription): return taskDescription.redactedDescription
-    case .preparation(let taskDescription): return taskDescription.redactedDescription
-    }
+    return wrapped.redactedDescription
   }
 
   public func execute() async {
-    switch self {
-    case .updateIndexStore(let taskDescription): return await taskDescription.execute()
-    case .preparation(let taskDescription): return await taskDescription.execute()
-    }
+    return await wrapped.execute()
   }
 
   /// Forward to the underlying task to compute the dependencies. Preparation and index tasks don't have any
   /// dependencies that are managed by `TaskScheduler`. `SemanticIndexManager` awaits the preparation of a target before
   /// indexing files within it.
   public func dependencies(
-    to currentlyExecutingTasks: [IndexTaskDescription]
-  ) -> [TaskDependencyAction<IndexTaskDescription>] {
-    switch self {
-    case .updateIndexStore(let taskDescription):
-      let currentlyExecutingTasks =
-        currentlyExecutingTasks
-        .compactMap { (currentlyExecutingTask) -> UpdateIndexStoreTaskDescription? in
-          if case .updateIndexStore(let currentlyExecutingTask) = currentlyExecutingTask {
-            return currentlyExecutingTask
-          }
-          return nil
-        }
-      return taskDescription.dependencies(to: currentlyExecutingTasks).map {
-        switch $0 {
-        case .waitAndElevatePriorityOfDependency(let td):
-          return .waitAndElevatePriorityOfDependency(.updateIndexStore(td))
-        case .cancelAndRescheduleDependency(let td):
-          return .cancelAndRescheduleDependency(.updateIndexStore(td))
-        }
-      }
-    case .preparation(let taskDescription):
-      let currentlyExecutingTasks =
-        currentlyExecutingTasks
-        .compactMap { (currentlyExecutingTask) -> PreparationTaskDescription? in
-          if case .preparation(let currentlyExecutingTask) = currentlyExecutingTask {
-            return currentlyExecutingTask
-          }
-          return nil
-        }
-      return taskDescription.dependencies(to: currentlyExecutingTasks).map {
-        switch $0 {
-        case .waitAndElevatePriorityOfDependency(let td):
-          return .waitAndElevatePriorityOfDependency(.preparation(td))
-        case .cancelAndRescheduleDependency(let td):
-          return .cancelAndRescheduleDependency(.preparation(td))
-        }
-      }
-    }
+    to currentlyExecutingTasks: [AnyIndexTaskDescription]
+  ) -> [TaskDependencyAction<AnyIndexTaskDescription>] {
+    return wrapped.dependencies(to: currentlyExecutingTasks)
   }
 }

Sources/SemanticIndex/PreparationTaskDescription.swift

@@ -24,7 +24,9 @@ private var preparationIDForLogging = AtomicUInt32(initialValue: 1)
 /// Describes a task to prepare a set of targets.
 ///
 /// This task description can be scheduled in a `TaskScheduler`.
-public struct PreparationTaskDescription: TaskDescriptionProtocol {
+public struct PreparationTaskDescription: IndexTaskDescription {
+  public static let idPrefix = "prepare"
+
   public let id = preparationIDForLogging.fetchAndIncrement()
 
   /// The targets that should be prepared.
@@ -33,11 +35,6 @@ public struct PreparationTaskDescription: TaskDescriptionProtocol {
   /// The build system manager that is used to get the toolchain and build settings for the files to index.
   private let buildSystemManager: BuildSystemManager
 
-  /// A callback that is called when the task finishes.
-  ///
-  /// Intended for testing purposes.
-  private let didFinishCallback: @Sendable (PreparationTaskDescription) -> Void
-
   /// The task is idempotent because preparing the same target twice produces the same result as preparing it once.
   public var isIdempotent: Bool { true }
 
@@ -53,18 +50,13 @@ public struct PreparationTaskDescription: TaskDescriptionProtocol {
 
   init(
     targetsToPrepare: [ConfiguredTarget],
-    buildSystemManager: BuildSystemManager,
-    didFinishCallback: @escaping @Sendable (PreparationTaskDescription) -> Void
+    buildSystemManager: BuildSystemManager
   ) {
     self.targetsToPrepare = targetsToPrepare
     self.buildSystemManager = buildSystemManager
-    self.didFinishCallback = didFinishCallback
   }
 
   public func execute() async {
-    defer {
-      didFinishCallback(self)
-    }
     // Only use the last two digits of the preparation ID for the logging scope to avoid creating too many scopes.
     // See comment in `withLoggingScope`.
     // The last 2 digits should be sufficient to differentiate between multiple concurrently running preparation operations