Make IndexTaskDescription protocol-based instead of enum-based

This simplifies the implementation.

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.

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -44,20 +44,20 @@ public final actor SemanticIndexManager {
   /// The `TaskScheduler` that manages the scheduling of index tasks. This is shared among all `SemanticIndexManager`s
   /// in the process, to ensure that we don't schedule more index operations than processor cores from multiple
   /// workspaces.
-  private let indexTaskScheduler: TaskScheduler<IndexTaskDescription>
+  private let indexTaskScheduler: TaskScheduler<AnyIndexTaskDescription>
 
   /// Callback that is called when an index task has finished.
   ///
   /// Currently only used for testing.
-  private let indexTaskDidFinish: (@Sendable (IndexTaskDescription) -> Void)?
+  private let indexTaskDidFinish: (@Sendable (AnyIndexTaskDescription) -> Void)?
 
   // MARK: - Public API
 
   public init(
     index: UncheckedIndex,
     buildSystemManager: BuildSystemManager,
-    indexTaskScheduler: TaskScheduler<IndexTaskDescription>,
-    indexTaskDidFinish: (@Sendable (IndexTaskDescription) -> Void)?
+    indexTaskScheduler: TaskScheduler<AnyIndexTaskDescription>,
+    indexTaskDidFinish: (@Sendable (AnyIndexTaskDescription) -> Void)?
   ) {
     self.index = index.checked(for: .modifiedFiles)
     self.buildSystemManager = buildSystemManager
@@ -133,12 +133,12 @@ public final actor SemanticIndexManager {
   private func prepare(targets: [ConfiguredTarget], priority: TaskPriority?) async {
     await self.indexTaskScheduler.schedule(
       priority: priority,
-      .preparation(
+      AnyIndexTaskDescription(
         PreparationTaskDescription(
           targetsToPrepare: targets,
           buildSystemManager: self.buildSystemManager,
           didFinishCallback: { [weak self] taskDescription in
-            self?.indexTaskDidFinish?(.preparation(taskDescription))
+            self?.indexTaskDidFinish?(AnyIndexTaskDescription(taskDescription))
           }
         )
       )
@@ -149,13 +149,13 @@ public final actor SemanticIndexManager {
   private func updateIndexStore(for files: [DocumentURI], priority: TaskPriority?) async {
     await self.indexTaskScheduler.schedule(
       priority: priority,
-      .updateIndexStore(
+      AnyIndexTaskDescription(
         UpdateIndexStoreTaskDescription(
           filesToIndex: Set(files),
           buildSystemManager: self.buildSystemManager,
           index: self.index.unchecked,
           didFinishCallback: { [weak self] taskDescription in
-            self?.indexTaskDidFinish?(.updateIndexStore(taskDescription))
+            self?.indexTaskDidFinish?(AnyIndexTaskDescription(taskDescription))
           }
         )
       )

Sources/SemanticIndex/UpdateIndexStoreTaskDescription.swift

@@ -25,7 +25,8 @@ private nonisolated(unsafe) var updateIndexStoreIDForLogging = AtomicUInt32(init
 /// Describes a task to index a set of source files.
 ///
 /// This task description can be scheduled in a `TaskScheduler`.
-public struct UpdateIndexStoreTaskDescription: TaskDescriptionProtocol {
+public struct UpdateIndexStoreTaskDescription: IndexTaskDescription {
+  public static let idPrefix = "update-indexstore"
   public let id = updateIndexStoreIDForLogging.fetchAndIncrement()
 
   /// The files that should be indexed.

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -454,7 +454,7 @@ public actor SourceKitLSPServer {
   ///
   /// Shared process-wide to ensure the scheduled index operations across multiple workspaces don't exceed the maximum
   /// number of processor cores that the user allocated to background indexing.
-  private let indexTaskScheduler: TaskScheduler<IndexTaskDescription>
+  private let indexTaskScheduler: TaskScheduler<AnyIndexTaskDescription>
 
   private var packageLoadingWorkDoneProgress = WorkDoneProgressState(
     "SourceKitLSP.SourceKitLSPServer.reloadPackage",

Sources/SourceKitLSP/Workspace.swift

@@ -89,7 +89,7 @@ public final class Workspace: Sendable {
     underlyingBuildSystem: BuildSystem?,
     index uncheckedIndex: UncheckedIndex?,
     indexDelegate: SourceKitIndexDelegate?,
-    indexTaskScheduler: TaskScheduler<IndexTaskDescription>
+    indexTaskScheduler: TaskScheduler<AnyIndexTaskDescription>
   ) async {
     self.documentManager = documentManager
     self.buildSetup = options.buildSetup
@@ -142,7 +142,7 @@ public final class Workspace: Sendable {
     options: SourceKitLSPServer.Options,
     compilationDatabaseSearchPaths: [RelativePath],
     indexOptions: IndexOptions = IndexOptions(),
-    indexTaskScheduler: TaskScheduler<IndexTaskDescription>,
+    indexTaskScheduler: TaskScheduler<AnyIndexTaskDescription>,
     reloadPackageStatusCallback: @Sendable @escaping (ReloadPackageStatus) async -> Void
   ) async throws {
     var buildSystem: BuildSystem? = nil
@@ -306,7 +306,7 @@ public struct IndexOptions: Sendable {
   /// A callback that is called when an index task finishes.
   ///
   /// Intended for testing purposes.
-  public var indexTaskDidFinish: (@Sendable (IndexTaskDescription) -> Void)?
+  public var indexTaskDidFinish: (@Sendable (AnyIndexTaskDescription) -> Void)?
 
   public init(
     indexStorePath: AbsolutePath? = nil,
@@ -315,7 +315,7 @@ public struct IndexOptions: Sendable {
     listenToUnitEvents: Bool = true,
     enableBackgroundIndexing: Bool = false,
     maxCoresPercentageToUseForBackgroundIndexing: Double = 1,
-    indexTaskDidFinish: (@Sendable (IndexTaskDescription) -> Void)? = nil
+    indexTaskDidFinish: (@Sendable (AnyIndexTaskDescription) -> Void)? = nil
   ) {
     self.indexStorePath = indexStorePath
     self.indexDatabasePath = indexDatabasePath