Support background preparation of targets

Author: ahoppen

Sources/SKCore/BuildServerBuildSystem.swift

@@ -279,6 +279,16 @@ extension BuildServerBuildSystem: BuildSystem {
     return [ConfiguredTarget(targetID: "dummy", runDestinationID: "dummy")]
   }
 
+  public func generateBuildGraph() {}
+
+  public func topologicalSort(of targets: [ConfiguredTarget]) async -> [ConfiguredTarget]? {
+    return nil
+  }
+
+  public func prepare(targets: [ConfiguredTarget]) async throws {
+    throw PrepareNotSupportedError()
+  }
+
   public func registerForChangeNotifications(for uri: DocumentURI) {
     let request = RegisterForChanges(uri: uri, action: .register)
     _ = self.buildServer?.send(request) { result in

Sources/SKCore/BuildSystem.swift

@@ -31,14 +31,19 @@ public struct SourceFileInfo: Sendable {
   /// The URI of the source file.
   public let uri: DocumentURI
 
+  /// `true` if this file belongs to the root project that the user is working on. It is false, if the file belongs
+  /// to a dependency of the project.
+  public let isPartOfRootProject: Bool
+
   /// Whether the file might contain test cases. This property is an over-approximation. It might be true for files
   /// from non-test targets or files that don't actually contain any tests. Keeping this list of files with
   /// `mayContainTets` minimal as possible helps reduce the amount of work that the syntactic test indexer needs to
   /// perform.
   public let mayContainTests: Bool
 
-  public init(uri: DocumentURI, mayContainTests: Bool) {
+  public init(uri: DocumentURI, isPartOfRootProject: Bool, mayContainTests: Bool) {
     self.uri = uri
+    self.isPartOfRootProject = isPartOfRootProject
     self.mayContainTests = mayContainTests
   }
 }
@@ -64,6 +69,13 @@ public struct ConfiguredTarget: Hashable, Sendable {
   }
 }
 
+/// An error build systems can throw from `prepare` if they don't support preparation of targets.
+public struct PrepareNotSupportedError: Error, CustomStringConvertible {
+  public init() {}
+
+  public var description: String { "Preparation not supported" }
+}
+
 /// Provider of FileBuildSettings and other build-related information.
 ///
 /// The primary role of the build system is to answer queries for
@@ -114,6 +126,22 @@ public protocol BuildSystem: AnyObject, Sendable {
   /// Return the list of targets and run destinations that the given document can be built for.
   func configuredTargets(for document: DocumentURI) async -> [ConfiguredTarget]
 
+  /// Re-generate the build graph including all the tasks that are necessary for building the entire build graph, like
+  /// resolving package versions.
+  func generateBuildGraph() async throws
+
+  /// Sort the targets so that low-level targets occur before high-level targets.
+  ///
+  /// This sorting is best effort but allows the indexer to prepare and index low-level targets first, which allows
+  /// index data to be available earlier.
+  ///
+  /// `nil` if the build system doesn't support topological sorting of targets.
+  func topologicalSort(of targets: [ConfiguredTarget]) async -> [ConfiguredTarget]?
+
+  /// Prepare the given targets for indexing and semantic functionality. This should build all swift modules of target
+  /// dependencies.
+  func prepare(targets: [ConfiguredTarget]) async throws
+
   /// If the build system has knowledge about the language that this document should be compiled in, return it.
   ///
   /// This is used to determine the language in which a source file should be background indexed.
@@ -146,5 +174,3 @@ public protocol BuildSystem: AnyObject, Sendable {
   /// The callback might also be called without an actual change to `sourceFiles`.
   func addSourceFilesDidChangeCallback(_ callback: @Sendable @escaping () async -> Void) async
 }
-
-public let buildTargetsNotSupported = ResponseError.methodNotFound(BuildTargets.method)

Sources/SKCore/BuildSystemManager.swift

@@ -208,6 +208,18 @@ extension BuildSystemManager {
     return settings
   }
 
+  public func generateBuildGraph() async throws {
+    try await self.buildSystem?.generateBuildGraph()
+  }
+
+  public func topologicalSort(of targets: [ConfiguredTarget]) async throws -> [ConfiguredTarget] {
+    return await buildSystem?.topologicalSort(of: targets) ?? targets
+  }
+
+  public func prepare(targets: [ConfiguredTarget]) async throws {
+    try await buildSystem?.prepare(targets: targets)
+  }
+
   public func registerForChangeNotifications(for uri: DocumentURI, language: Language) async {
     logger.debug("registerForChangeNotifications(\(uri.forLogging))")
     let mainFile = await mainFile(for: uri, language: language)
@@ -247,7 +259,7 @@ extension BuildSystemManager {
 
   public func testFiles() async -> [DocumentURI] {
     return await sourceFiles().compactMap { (info: SourceFileInfo) -> DocumentURI? in
-      guard info.mayContainTests else {
+      guard info.isPartOfRootProject, info.mayContainTests else {
         return nil
       }
       return info.uri

Sources/SKCore/CompilationDatabaseBuildSystem.swift

@@ -125,6 +125,16 @@ extension CompilationDatabaseBuildSystem: BuildSystem {
     return [ConfiguredTarget(targetID: "dummy", runDestinationID: "dummy")]
   }
 
+  public func prepare(targets: [ConfiguredTarget]) async throws {
+    throw PrepareNotSupportedError()
+  }
+
+  public func generateBuildGraph() {}
+
+  public func topologicalSort(of targets: [ConfiguredTarget]) -> [ConfiguredTarget]? {
+    return nil
+  }
+
   public func registerForChangeNotifications(for uri: DocumentURI) async {
     self.watchedFiles.insert(uri)
   }
@@ -208,7 +218,7 @@ extension CompilationDatabaseBuildSystem: BuildSystem {
       return []
     }
     return compdb.allCommands.map {
-      SourceFileInfo(uri: DocumentURI($0.url), mayContainTests: true)
+      SourceFileInfo(uri: DocumentURI($0.url), isPartOfRootProject: true, mayContainTests: true)
     }
   }
 

Sources/SKSupport/Collection+PartitionIntoBatches.swift

@@ -10,7 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-public extension Collection {
+public extension Collection where Index == Int {
   /// Partition the elements of the collection into `numberOfBatches` roughly equally sized batches.
   ///
   /// Elements are assigned to the batches round-robin. This ensures that elements that are close to each other in the
@@ -32,4 +32,18 @@ public extension Collection {
     }
     return batches.filter { !$0.isEmpty }
   }
+
+  /// Partition the collection into batches that have a maximum size of `batchSize`.
+  ///
+  /// The last batch will contain the remainder elements.
+  func partition(intoBatchesOfSize batchSize: Int) -> [[Element]] {
+    var batches: [[Element]] = []
+    batches.reserveCapacity(self.count / batchSize)
+    var lastIndex = self.startIndex
+    for index in stride(from: self.startIndex, to: self.endIndex, by: batchSize).dropFirst() + [self.endIndex] {
+      batches.append(Array(self[lastIndex..<index]))
+      lastIndex = index
+    }
+    return batches
+  }
 }

Sources/SKSwiftPMWorkspace/SwiftPMBuildSystem.swift

@@ -237,6 +237,16 @@ public actor SwiftPMBuildSystem {
 }
 
 extension SwiftPMBuildSystem {
+  public func generateBuildGraph() async throws {
+    let observabilitySystem = ObservabilitySystem({ scope, diagnostic in
+      logger.log(level: diagnostic.severity.asLogLevel, "SwiftPM log: \(diagnostic.description)")
+    })
+    try self.workspace.resolve(
+      root: PackageGraphRootInput(packages: [AbsolutePath(projectRoot)]),
+      observabilityScope: observabilitySystem.topScope
+    )
+    try await self.reloadPackage()
+  }
 
   /// (Re-)load the package settings by parsing the manifest and resolving all the targets and
   /// dependencies.
@@ -395,6 +405,72 @@ extension SwiftPMBuildSystem: SKCore.BuildSystem {
     return []
   }
 
+  public func topologicalSort(of targets: [ConfiguredTarget]) -> [ConfiguredTarget]? {
+    return targets.sorted { (lhs: ConfiguredTarget, rhs: ConfiguredTarget) -> Bool in
+      let lhsIndex = self.targets.firstIndex(where: { $0.name == lhs.targetID }) ?? self.targets.count
+      let rhsIndex = self.targets.firstIndex(where: { $0.name == rhs.targetID }) ?? self.targets.count
+      return lhsIndex < rhsIndex
+    }
+  }
+
+  public func prepare(targets: [ConfiguredTarget]) async throws {
+    // TODO (indexing): Support preparation of multiple targets at once.
+    for target in targets {
+      try await prepare(singleTarget: target)
+    }
+  }
+
+  private func prepare(singleTarget target: ConfiguredTarget) async throws {
+    // TODO (indexing): Add a proper 'prepare' jobs in SwiftPM instead of building the target
+    guard let toolchain = await toolchainRegistry.default else {
+      logger.error("Not preparing because not toolchain exists")
+      return
+    }
+    guard let swift = toolchain.swift else {
+      logger.error(
+        "Not preparing because toolchain at \(toolchain.identifier) does not contain a Swift compiler"
+      )
+      return
+    }
+    let arguments = [
+      swift.pathString, "build",
+      "--scratch-path", self.workspace.location.scratchDirectory.pathString,
+      "--disable-index-store",
+      "--target", target.targetID,
+    ]
+    let process = Process(
+      arguments: arguments,
+      workingDirectory: try TSCBasic.AbsolutePath(validating: workspacePath.pathString)
+    )
+    try process.launch()
+    let result = try await process.waitUntilExitSendingSigIntOnTaskCancellation()
+    switch result.exitStatus {
+    case .terminated(code: 0):
+      break
+    case .terminated(code: let code):
+      // This most likely happens if there are compilation errors in the source file. This is nothing to worry about.
+      let stdout = (try? String(bytes: result.output.get(), encoding: .utf8)) ?? "<no stderr>"
+      let stderr = (try? String(bytes: result.stderrOutput.get(), encoding: .utf8)) ?? "<no stderr>"
+      logger.debug(
+        """
+        Preparation of targets \(target.targetID) terminated with non-zero exit code \(code)
+        Stderr:
+        \(stderr)
+        Stdout:
+        \(stdout)
+        """
+      )
+    // await BuildSettingsLogger.index.log(settings: buildSettings, for: uri)
+    case .signalled(signal: let signal):
+      if !Task.isCancelled {
+        // The indexing job finished with a signal. Could be because the compiler crashed.
+        // Ignore signal exit codes if this task has been cancelled because the compiler exits with SIGINT if it gets
+        // interrupted.
+        logger.error("Preparation of targets \(target.targetID) signaled \(signal)")
+      }
+    }
+  }
+
   public func registerForChangeNotifications(for uri: DocumentURI) async {
     self.watchedFiles.insert(uri)
   }
@@ -489,14 +565,11 @@ extension SwiftPMBuildSystem: SKCore.BuildSystem {
 
   public func sourceFiles() -> [SourceFileInfo] {
     return fileToTarget.compactMap { (path, target) -> SourceFileInfo? in
-      guard target.isPartOfRootPackage else {
-        // Don't consider files from package dependencies as possible test files.
-        return nil
-      }
       // We should only set mayContainTests to `true` for files from test targets
       // (https://github.com/apple/sourcekit-lsp/issues/1174).
       return SourceFileInfo(
         uri: DocumentURI(path.asURL),
+        isPartOfRootProject: target.isPartOfRootPackage,
         mayContainTests: true
       )
     }

Sources/SKTestSupport/FileManager+findFiles.swift

@@ -24,4 +24,16 @@ extension FileManager {
     }
     return result
   }
+
+  /// Returns the URLs of all files with the given file extension in the given directory (recursively).
+  public func findFiles(named name: String, in directory: URL) -> [URL] {
+    var result: [URL] = []
+    let enumerator = self.enumerator(at: directory, includingPropertiesForKeys: nil)
+    while let url = enumerator?.nextObject() as? URL {
+      if url.lastPathComponent == name {
+        result.append(url)
+      }
+    }
+    return result
+  }
 }

Sources/SKTestSupport/SwiftPMDependencyProject.swift

@@ -71,13 +71,13 @@ public class SwiftPMDependencyProject {
     var files = files
     files["Package.swift"] = manifest
 
-    for (fileLocation, contents) in files {
+    for (fileLocation, markedContents) in files {
       let fileURL = fileLocation.url(relativeTo: packageDirectory)
       try FileManager.default.createDirectory(
         at: fileURL.deletingLastPathComponent(),
         withIntermediateDirectories: true
       )
-      try contents.write(to: fileURL, atomically: true, encoding: .utf8)
+      try extractMarkers(markedContents).textWithoutMarkers.write(to: fileURL, atomically: true, encoding: .utf8)
     }
 
     try await runGitCommand(["init"], workingDirectory: packageDirectory)

Sources/SemanticIndex/CMakeLists.txt

@@ -1,6 +1,8 @@
 
 add_library(SemanticIndex STATIC
   CheckedIndex.swift
+  IndexTaskDescription.swift
+  PreparationTaskDescription.swift
   SemanticIndexManager.swift
   UpdateIndexStoreTaskDescription.swift
 )