Make Workspace the delegate of a BuildSystemManager

`Workspace` is responsible for creating the `BuildSystemManager` and responds to most of the delegate calls. It should thus also be the delegate of `BuildSystemManager`.

Author: ahoppen

Sources/BuildSystemIntegration/BuildSystemManager.swift

@@ -87,7 +87,7 @@ package actor BuildSystemManager: BuiltInBuildSystemAdapterDelegate {
   var mainFilesProvider: MainFilesProvider?
 
   /// Build system delegate that will receive notifications about setting changes, etc.
-  var delegate: BuildSystemManagerDelegate?
+  private weak var delegate: BuildSystemManagerDelegate?
 
   /// The list of toolchains that are available.
   ///
@@ -157,6 +157,11 @@ package actor BuildSystemManager: BuiltInBuildSystemAdapterDelegate {
 
   package func filesDidChange(_ events: [FileEvent]) async {
     await self.buildSystem?.send(BuildServerProtocol.DidChangeWatchedFilesNotification(changes: events))
+    if let mainFilesProvider {
+      var mainFiles = await Set(events.asyncFlatMap { await mainFilesProvider.mainFilesContainingFile($0.uri) })
+      mainFiles.subtract(events.map(\.uri))
+      await self.filesDependenciesUpdated(mainFiles)
+    }
   }
 
   /// Implementation of `MessageHandler`, handling notifications from the build system.

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -842,67 +842,6 @@ extension SourceKitLSPServer: MessageHandler {
   }
 }
 
-// MARK: - Build System Delegate
-
-extension SourceKitLSPServer: BuildSystemManagerDelegate {
-  private func affectedOpenDocumentsForChangeSet(
-    _ changes: Set<DocumentURI>,
-    _ documentManager: DocumentManager
-  ) -> Set<DocumentURI> {
-    // An empty change set is treated as if all open files have been modified.
-    guard !changes.isEmpty else {
-      return documentManager.openDocuments
-    }
-    return documentManager.openDocuments.intersection(changes)
-  }
-
-  /// Handle a build settings change notification from the `BuildSystem`.
-  /// This has two primary cases:
-  /// - Initial settings reported for a given file, now we can fully open it
-  /// - Changed settings for an already open file
-  package func fileBuildSettingsChanged(_ changedFiles: Set<DocumentURI>) async {
-    for uri in changedFiles {
-      guard self.documentManager.openDocuments.contains(uri) else {
-        continue
-      }
-
-      guard let service = await self.workspaceForDocument(uri: uri)?.documentService(for: uri) else {
-        continue
-      }
-
-      await service.documentUpdatedBuildSettings(uri)
-    }
-  }
-
-  /// Handle a dependencies updated notification from the `BuildSystem`.
-  /// We inform the respective language services as long as the given file is open
-  /// (not queued for opening).
-  package func filesDependenciesUpdated(_ changedFiles: Set<DocumentURI>) async {
-    // Split the changedFiles into the workspaces they belong to.
-    // Then invoke affectedOpenDocumentsForChangeSet for each workspace with its affected files.
-    let changedFilesAndWorkspace = await changedFiles.asyncMap {
-      return (uri: $0, workspace: await self.workspaceForDocument(uri: $0))
-    }
-    for workspace in self.workspaces {
-      let changedFilesForWorkspace = Set(changedFilesAndWorkspace.filter({ $0.workspace === workspace }).map(\.uri))
-      if changedFilesForWorkspace.isEmpty {
-        continue
-      }
-      for uri in self.affectedOpenDocumentsForChangeSet(changedFilesForWorkspace, self.documentManager) {
-        logger.log("Dependencies updated for opened file \(uri.forLogging)")
-        if let service = workspace.documentService(for: uri) {
-          await service.documentDependenciesUpdated(uri)
-        }
-      }
-    }
-  }
-
-  package func fileHandlingCapabilityChanged() {
-    logger.log("Updating URI to workspace because file handling capability of a workspace changed")
-    self.scheduleUpdateOfUriToWorkspace()
-  }
-}
-
 extension SourceKitLSPServer {
   nonisolated func logMessageToIndexLog(taskID: IndexTaskID, message: String) {
     var message: Substring = message[...]
@@ -980,6 +919,10 @@ extension SourceKitLSPServer {
       },
       reloadPackageStatusCallback: { [weak self] status in
         await self?.reloadPackageStatusCallback(status)
+      },
+      fileHandlingCapabilityChanged: { [weak self] in
+        logger.log("Updating URI to workspace because file handling capability of a workspace changed")
+        await self?.scheduleUpdateOfUriToWorkspace()
       }
     )
     if options.backgroundIndexingOrDefault, workspace.semanticIndexManager == nil,
@@ -1111,6 +1054,10 @@ extension SourceKitLSPServer {
           },
           reloadPackageStatusCallback: { [weak self] status in
             await self?.reloadPackageStatusCallback(status)
+          },
+          fileHandlingCapabilityChanged: { [weak self] in
+            logger.log("Updating URI to workspace because file handling capability of a workspace changed")
+            await self?.scheduleUpdateOfUriToWorkspace()
           }
         )
 
@@ -1119,9 +1066,6 @@ extension SourceKitLSPServer {
     }.value
 
     assert(!self.workspaces.isEmpty)
-    for workspace in self.workspaces {
-      await workspace.buildSystemManager.setDelegate(self)
-    }
 
     return InitializeResult(
       capabilities: await self.serverCapabilities(
@@ -1340,9 +1284,6 @@ extension SourceKitLSPServer {
     for workspace in self.workspaces {
       await workspace.buildSystemManager.setMainFilesProvider(nil)
       workspace.closeIndex()
-
-      // Break retain cycle with the BSM.
-      await workspace.buildSystemManager.setDelegate(nil)
     }
   }
 
@@ -1541,9 +1482,6 @@ extension SourceKitLSPServer {
             )
           }
         }
-        for workspace in newWorkspaces {
-          await workspace.buildSystemManager.setDelegate(self)
-        }
         self.workspacesAndIsImplicit += newWorkspaces.map { (workspace: $0, isImplicit: false) }
       }
     }.value
@@ -1940,7 +1878,7 @@ extension SourceKitLSPServer {
         == canonicalOriginatorLocation
     }
 
-    var locations = try await symbols.asyncMap { (symbol) -> [Location] in
+    var locations = try await symbols.asyncFlatMap { (symbol) -> [Location] in
       var locations: [Location]
       if let bestLocalDeclaration = symbol.bestLocalDeclaration,
         !(symbol.isDynamic ?? true),
@@ -1978,7 +1916,7 @@ extension SourceKitLSPServer {
       }
 
       return locations
-    }.flatMap { $0 }
+    }
 
     // Remove any duplicate locations. We might end up with duplicate locations when performing a definition request
     // on eg. `MyStruct()` when no explicit initializer is declared. In this case we get two symbol infos, one for the

Sources/SourceKitLSP/Workspace.swift

@@ -48,7 +48,7 @@ fileprivate func firstNonNil<T>(
 /// "initialize" request has been made.
 ///
 /// Typically a workspace is contained in a root directory.
-package final class Workspace: Sendable {
+package final class Workspace: Sendable, BuildSystemManagerDelegate {
 
   /// The root directory of the workspace.
   package let rootUri: DocumentURI?
@@ -85,6 +85,9 @@ package final class Workspace: Sendable {
   /// `nil` if background indexing is not enabled.
   let semanticIndexManager: SemanticIndexManager?
 
+  /// A callback that should be called when the file handling capability of this workspace changes.
+  private let fileHandlingCapabilityChangedCallback: @Sendable () async -> Void
+
   private init(
     documentManager: DocumentManager,
     rootUri: DocumentURI?,
@@ -97,14 +100,16 @@ package final class Workspace: Sendable {
     indexTaskScheduler: TaskScheduler<AnyIndexTaskDescription>,
     logMessageToIndexLog: @escaping @Sendable (_ taskID: IndexTaskID, _ message: String) -> Void,
     indexTasksWereScheduled: @escaping @Sendable (Int) -> Void,
-    indexProgressStatusDidChange: @escaping @Sendable () -> Void
+    indexProgressStatusDidChange: @escaping @Sendable () -> Void,
+    fileHandlingCapabilityChanged: @escaping @Sendable () async -> Void
   ) async {
     self.documentManager = documentManager
     self.rootUri = rootUri
     self.capabilityRegistry = capabilityRegistry
     self.options = options
     self._uncheckedIndex = ThreadSafeBox(initialValue: uncheckedIndex)
     self.buildSystemManager = buildSystemManager
+    self.fileHandlingCapabilityChangedCallback = fileHandlingCapabilityChanged
     if options.backgroundIndexingOrDefault, let uncheckedIndex, await buildSystemManager.supportsPreparation {
       self.semanticIndexManager = SemanticIndexManager(
         index: uncheckedIndex,
@@ -153,7 +158,8 @@ package final class Workspace: Sendable {
     logMessageToIndexLog: @escaping @Sendable (_ taskID: IndexTaskID, _ message: String) -> Void,
     indexTasksWereScheduled: @Sendable @escaping (Int) -> Void,
     indexProgressStatusDidChange: @Sendable @escaping () -> Void,
-    reloadPackageStatusCallback: @Sendable @escaping (ReloadPackageStatus) async -> Void
+    reloadPackageStatusCallback: @Sendable @escaping (ReloadPackageStatus) async -> Void,
+    fileHandlingCapabilityChanged: @Sendable @escaping () async -> Void
   ) async {
     let buildSystemManager = await BuildSystemManager(
       buildSystemKind: buildSystemKind,
@@ -227,8 +233,10 @@ package final class Workspace: Sendable {
       indexTaskScheduler: indexTaskScheduler,
       logMessageToIndexLog: logMessageToIndexLog,
       indexTasksWereScheduled: indexTasksWereScheduled,
-      indexProgressStatusDidChange: indexProgressStatusDidChange
+      indexProgressStatusDidChange: indexProgressStatusDidChange,
+      fileHandlingCapabilityChanged: fileHandlingCapabilityChanged
     )
+    await buildSystemManager.setDelegate(self)
   }
 
   package static func forTesting(
@@ -249,7 +257,8 @@ package final class Workspace: Sendable {
       indexTaskScheduler: indexTaskScheduler,
       logMessageToIndexLog: { _, _ in },
       indexTasksWereScheduled: { _ in },
-      indexProgressStatusDidChange: {}
+      indexProgressStatusDidChange: {},
+      fileHandlingCapabilityChanged: {}
     )
   }
 
@@ -291,6 +300,39 @@ package final class Workspace: Sendable {
       return newLanguageService
     }
   }
+
+  /// Handle a build settings change notification from the `BuildSystem`.
+  /// This has two primary cases:
+  /// - Initial settings reported for a given file, now we can fully open it
+  /// - Changed settings for an already open file
+  package func fileBuildSettingsChanged(_ changedFiles: Set<DocumentURI>) async {
+    for uri in changedFiles {
+      guard self.documentManager.openDocuments.contains(uri) else {
+        continue
+      }
+
+      await self.documentService(for: uri)?.documentUpdatedBuildSettings(uri)
+    }
+  }
+
+  /// Handle a dependencies updated notification from the `BuildSystem`.
+  /// We inform the respective language services as long as the given file is open
+  /// (not queued for opening).
+  package func filesDependenciesUpdated(_ changedFiles: Set<DocumentURI>) async {
+    for uri in changedFiles {
+      guard self.documentManager.openDocuments.contains(uri) else {
+        continue
+      }
+      logger.log("Dependencies updated for opened file \(uri.forLogging)")
+      if let service = documentService(for: uri) {
+        await service.documentDependenciesUpdated(uri)
+      }
+    }
+  }
+
+  package func fileHandlingCapabilityChanged() async {
+    await self.fileHandlingCapabilityChangedCallback()
+  }
 }
 
 /// Wrapper around a workspace that isn't being retained.

Sources/SwiftExtensions/Sequence+AsyncMap.swift

@@ -25,6 +25,20 @@ extension Sequence {
     return result
   }
 
+  /// Just like `Sequence.flatMap` but allows an `async` transform function.
+  package func asyncFlatMap<SegmentOfResult: Sequence>(
+    @_inheritActorContext _ transform: @Sendable (Element) async throws -> SegmentOfResult
+  ) async rethrows -> [SegmentOfResult.Element] {
+    var result: [SegmentOfResult.Element] = []
+    result.reserveCapacity(self.underestimatedCount)
+
+    for element in self {
+      result += try await transform(element)
+    }
+
+    return result
+  }
+
   /// Just like `Sequence.compactMap` but allows an `async` transform function.
   package func asyncCompactMap<T>(
     @_inheritActorContext _ transform: @Sendable (Element) async throws -> T?

Tests/SourceKitLSPTests/BuildSystemTests.swift

@@ -157,7 +157,7 @@ final class BuildSystemTests: XCTestCase {
     )
 
     await server.setWorkspaces([(workspace: workspace, isImplicit: false)])
-    await workspace.buildSystemManager.setDelegate(server)
+    await workspace.buildSystemManager.setDelegate(workspace)
   }
 
   override func tearDown() {

Tests/SourceKitLSPTests/DependencyTrackingTests.swift

@@ -40,7 +40,7 @@ final class DependencyTrackingTests: XCTestCase {
       usePullDiagnostics: false
     )
 
-    let (libBUri, _) = try project.openDocument("LibB.swift")
+    _ = try project.openDocument("LibB.swift")
 
     let initialDiags = try await project.testClient.nextDiagnosticsNotification()
     // Semantic analysis: expect module import error.
@@ -58,7 +58,14 @@ final class DependencyTrackingTests: XCTestCase {
 
     try await SwiftPMTestProject.build(at: project.scratchDirectory)
 
-    await project.testClient.server.filesDependenciesUpdated([libBUri])
+    project.testClient.send(
+      DidChangeWatchedFilesNotification(
+        changes:
+          FileManager.default.findFiles(withExtension: "swiftmodule", in: project.scratchDirectory).map {
+            FileEvent(uri: DocumentURI($0), type: .created)
+          }
+      )
+    )
 
     let updatedDiags = try await project.testClient.nextDiagnosticsNotification()
     // Semantic analysis: no more errors expected, import should resolve since we built.
@@ -102,7 +109,8 @@ final class DependencyTrackingTests: XCTestCase {
     let contents = "int libX(int value);"
     try contents.write(to: generatedHeaderURL, atomically: true, encoding: .utf8)
 
-    await project.testClient.server.filesDependenciesUpdated([mainUri])
+    let workspace = try await unwrap(project.testClient.server.workspaceForDocument(uri: mainUri))
+    await workspace.filesDependenciesUpdated([mainUri])
 
     let updatedDiags = try await project.testClient.nextDiagnosticsNotification()
     // No more errors expected, import should resolve since we the generated header file