Implement initial background indexing of a project

.gitignore

@@ -2,6 +2,7 @@
 default.profraw
 Package.resolved
 /.build
+/.index-build
 /Packages
 /*.xcodeproj
 /*.sublime-project

Package.swift

@@ -173,6 +173,8 @@ let package = Package(
     .target(
       name: "SemanticIndex",
       dependencies: [
+        "CAtomics",
+        "LanguageServerProtocol",
         "LSPLogging",
         "SKCore",
         .product(name: "IndexStoreDB", package: "indexstore-db"),

Sources/LSPLogging/Logging.swift

@@ -13,7 +13,7 @@
 /// Which log level to use (from https://developer.apple.com/wwdc20/10168?time=604)
 ///  - Debug: Useful only during debugging (only logged during debugging)
 ///  - Info: Helpful but not essential for troubleshooting (not persisted, logged to memory)
-///  - Notice/log (Default): Essential for troubleshooting
+///  - Notice/log/default: Essential for troubleshooting
 ///  - Error: Error seen during execution
 ///    - Used eg. if the user sends an erroneous request or if a request fails
 ///  - Fault: Bug in program

Sources/LSPLogging/NonDarwinLogging.swift

@@ -330,7 +330,7 @@ public struct NonDarwinLogger: Sendable {
     log(level: .info, message)
   }
 
-  /// Log a message at the `log` level.
+  /// Log a message at the `default` level.
   public func log(_ message: NonDarwinLogMessage) {
     log(level: .default, message)
   }

Sources/SKCore/BuildSystemManager.swift

@@ -159,7 +159,8 @@ extension BuildSystemManager {
   /// references to that C file in the build settings by the header file.
   public func buildSettingsInferredFromMainFile(
     for document: DocumentURI,
-    language: Language
+    language: Language,
+    logBuildSettings: Bool = true
   ) async -> FileBuildSettings? {
     let mainFile = await mainFile(for: document, language: language)
     guard var settings = await buildSettings(for: mainFile, language: language) else {
@@ -170,7 +171,9 @@ extension BuildSystemManager {
       // to reference `document` instead of `mainFile`.
       settings = settings.patching(newFile: document.pseudoPath, originalFile: mainFile.pseudoPath)
     }
-    await BuildSettingsLogger.shared.log(settings: settings, for: document)
+    if logBuildSettings {
+      await BuildSettingsLogger.shared.log(settings: settings, for: document)
+    }
     return settings
   }
 
@@ -349,16 +352,24 @@ extension BuildSystemManager {
 // MARK: - Build settings logger
 
 /// Shared logger that only logs build settings for a file once unless they change
-fileprivate actor BuildSettingsLogger {
-  static let shared = BuildSettingsLogger()
+public actor BuildSettingsLogger {
+  public static let shared = BuildSettingsLogger()
 
   private var loggedSettings: [DocumentURI: FileBuildSettings] = [:]
 
-  func log(settings: FileBuildSettings, for uri: DocumentURI) {
+  public func log(level: LogLevel = .default, settings: FileBuildSettings, for uri: DocumentURI) {
     guard loggedSettings[uri] != settings else {
       return
     }
     loggedSettings[uri] = settings
+    Self.log(level: level, settings: settings, for: uri)
+  }
+
+  /// Log the given build settings.
+  ///
+  /// In contrast to the instance method `log`, this will always log the build settings. The instance method only logs
+  /// the build settings if they have changed.
+  public static func log(level: LogLevel = .default, settings: FileBuildSettings, for uri: DocumentURI) {
     let log = """
       Compiler Arguments:
       \(settings.compilerArguments.joined(separator: "\n"))
@@ -370,6 +381,7 @@ fileprivate actor BuildSettingsLogger {
     let chunks = splitLongMultilineMessage(message: log)
     for (index, chunk) in chunks.enumerated() {
       logger.log(
+        level: level,
         """
         Build settings for \(uri.forLogging) (\(index + 1)/\(chunks.count))
         \(chunk)

Sources/SKTestSupport/SwiftPMTestProject.swift

@@ -43,6 +43,7 @@ public class SwiftPMTestProject: MultiFileTestProject {
     build: Bool = false,
     allowBuildFailure: Bool = false,
     serverOptions: SourceKitLSPServer.Options = .testDefault,
+    pollIndex: Bool = true,
     usePullDiagnostics: Bool = true,
     testName: String = #function
   ) async throws {
@@ -77,8 +78,10 @@ public class SwiftPMTestProject: MultiFileTestProject {
         try await Self.build(at: self.scratchDirectory)
       }
     }
-    // Wait for the indexstore-db to finish indexing
-    _ = try await testClient.send(PollIndexRequest())
+    if pollIndex {
+      // Wait for the indexstore-db to finish indexing
+      _ = try await testClient.send(PollIndexRequest())
+    }
   }
 
   /// Build a SwiftPM package package manifest is located in the directory at `path`.

Sources/SKTestSupport/WrappedSemaphore.swift

@@ -0,0 +1,36 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import Dispatch
+
+/// Wrapper around `DispatchSemaphore` so that Swift Concurrency doesn't complain about the usage of semaphores in the
+/// tests.
+///
+/// This should only be used for tests that test priority escalation and thus cannot await a `Task` (which would cause
+/// priority elevations).
+public struct WrappedSemaphore {
+  let semaphore = DispatchSemaphore(value: 0)
+
+  public init() {}
+
+  public func signal(value: Int = 1) {
+    for _ in 0..<value {
+      semaphore.signal()
+    }
+  }
+
+  public func wait(value: Int = 1) {
+    for _ in 0..<value {
+      semaphore.wait()
+    }
+  }
+}

Sources/SemanticIndex/CMakeLists.txt

@@ -1,6 +1,8 @@
 
 add_library(SemanticIndex STATIC
   CheckedIndex.swift
+  SemanticIndexManager.swift
+  UpdateIndexStoreTaskDescription.swift
 )
 set_target_properties(SemanticIndex PROPERTIES
   INTERFACE_INCLUDE_DIRECTORIES ${CMAKE_Swift_MODULE_DIRECTORY})

Sources/SemanticIndex/SemanticIndexManager.swift

@@ -0,0 +1,167 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import Foundation
+import LSPLogging
+import LanguageServerProtocol
+import SKCore
+
+/// Describes the state of indexing for a single source file
+private enum FileIndexStatus {
+  /// The index is up-to-date.
+  case upToDate
+  /// The file is being indexed by the given task.
+  case inProgress(Task<Void, Never>)
+}
+
+/// 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
+  /// need to be indexed again.
+  private let index: CheckedIndex
+
+  /// The build system manager that is used to get compiler arguments for a file.
+  private let buildSystemManager: BuildSystemManager
+
+  /// The index status of the source files that the `SemanticIndexManager` knows about.
+  ///
+  /// Files that have never been indexed are not in this dictionary.
+  private var indexStatus: [DocumentURI: FileIndexStatus] = [:]
+
+  /// 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<UpdateIndexStoreTaskDescription>
+
+  /// Callback that is called when an index task has finished.
+  ///
+  /// Currently only used for testing.
+  private let indexTaskDidFinish: (@Sendable (UpdateIndexStoreTaskDescription) -> Void)?
+
+  // MARK: - Public API
+
+  public init(
+    index: UncheckedIndex,
+    buildSystemManager: BuildSystemManager,
+    indexTaskScheduler: TaskScheduler<UpdateIndexStoreTaskDescription>,
+    indexTaskDidFinish: (@Sendable (UpdateIndexStoreTaskDescription) -> Void)?
+  ) {
+    self.index = index.checked(for: .modifiedFiles)
+    self.buildSystemManager = buildSystemManager
+    self.indexTaskScheduler = indexTaskScheduler
+    self.indexTaskDidFinish = indexTaskDidFinish
+  }
+
+  /// Schedules a task to index all files in `files` that don't already have an up-to-date index.
+  /// Returns immediately after scheduling that task.
+  ///
+  /// Indexing is being performed with a low priority.
+  public func scheduleBackgroundIndex(files: some Collection<DocumentURI>) {
+    self.index(files: files, priority: .low)
+  }
+
+  /// Wait for all in-progress index tasks to finish.
+  public func waitForUpToDateIndex() async {
+    logger.info("Waiting for up-to-date index")
+    await withTaskGroup(of: Void.self) { taskGroup in
+      for (_, status) in indexStatus {
+        switch status {
+        case .inProgress(let task):
+          taskGroup.addTask {
+            await task.value
+          }
+        case .upToDate:
+          break
+        }
+      }
+      await taskGroup.waitForAll()
+    }
+    index.pollForUnitChangesAndWait()
+    logger.debug("Done waiting for up-to-date index")
+  }
+
+  /// Ensure that the index for the given files is up-to-date.
+  ///
+  /// This tries to produce an up-to-date index for the given files as quickly as possible. To achieve this, it might
+  /// suspend previous target-wide index tasks in favor of index tasks that index a fewer files.
+  public func waitForUpToDateIndex(for uris: some Collection<DocumentURI>) async {
+    logger.info(
+      "Waiting for up-to-date index for \(uris.map { $0.fileURL?.lastPathComponent ?? $0.stringValue }.joined(separator: ", "))"
+    )
+    let filesWithOutOfDateIndex = uris.filter { uri in
+      switch indexStatus[uri] {
+      case .inProgress, nil: return true
+      case .upToDate: return false
+      }
+    }
+    // Create a new index task for the files that aren't up-to-date. The newly scheduled index tasks will
+    // - Wait for the existing index operations to finish if they have the same number of files.
+    // - Reschedule the background index task in favor of an index task with fewer source files.
+    await self.index(files: filesWithOutOfDateIndex, priority: nil).value
+    index.pollForUnitChangesAndWait()
+    logger.debug("Done waiting for up-to-date index")
+  }
+
+  // MARK: - Helper functions
+
+  /// Index the given set of files at the given priority.
+  ///
+  /// The returned task finishes when all files are indexed.
+  @discardableResult
+  private func index(files: some Collection<DocumentURI>, priority: TaskPriority?) -> Task<Void, Never> {
+    let outOfDateFiles = files.filter {
+      if case .upToDate = indexStatus[$0] {
+        return false
+      }
+      return true
+    }
+
+    var indexTasks: [Task<Void, Never>] = []
+
+    // TODO (indexing): Group index operations by target when we support background preparation.
+    for files in outOfDateFiles.partition(intoNumberOfBatches: ProcessInfo.processInfo.processorCount * 5) {
+      let indexTask = Task(priority: priority) {
+        await self.indexTaskScheduler.schedule(
+          priority: priority,
+          UpdateIndexStoreTaskDescription(
+            filesToIndex: Set(files),
+            buildSystemManager: self.buildSystemManager,
+            index: self.index,
+            didFinishCallback: { [weak self] taskDescription in
+              self?.indexTaskDidFinish?(taskDescription)
+            }
+          )
+        ).value
+        for file in files {
+          indexStatus[file] = .upToDate
+        }
+      }
+      indexTasks.append(indexTask)
+
+      for file in files {
+        indexStatus[file] = .inProgress(indexTask)
+      }
+    }
+    let indexTasksImmutable = indexTasks
+
+    return Task(priority: priority) {
+      await withTaskGroup(of: Void.self) { taskGroup in
+        for indexTask in indexTasksImmutable {
+          taskGroup.addTask(priority: priority) {
+            await indexTask.value
+          }
+        }
+        await taskGroup.waitForAll()
+      }
+    }
+  }
+}