(WIP) Optimize compiler argument aqueries

Example/.bsp/config.json

@@ -33,6 +33,7 @@
 		"--build-test-suffix",
 		"_(PLAT)_skbsp",
 		"--build-test-platform-placeholder",
-		"(PLAT)"
+		"(PLAT)",
+        "--separate-aquery-output"
 	]
 }

Sources/SourceKitBazelBSP/RequestHandlers/BuildTargets/BazelTargetStore.swift

@@ -30,14 +30,20 @@ private let logger = makeFileLevelBSPLogger()
 // the project's dependency graph and its files.
 protocol BazelTargetStore: AnyObject {
     var stateLock: OSAllocatedUnfairLock<Void> { get }
+    var platformsToTopLevelLabelsMap: [String: [String]] { get }
     func fetchTargets() throws -> [BuildTarget]
     func bazelTargetLabel(forBSPURI uri: URI) throws -> String
     func bazelTargetSrcs(forBSPURI uri: URI) throws -> [URI]
     func bspURIs(containingSrc src: URI) throws -> [URI]
-    func platformBuildLabel(forBSPURI uri: URI) throws -> (String, TopLevelRuleType)
+    func platformBuildLabelInfo(forBSPURI uri: URI) throws -> BazelTargetPlatformInfo
     func clearCache()
 }
 
+struct BazelTargetPlatformInfo {
+    let buildTestLabel: String
+    let parentRuleType: TopLevelRuleType
+}
+
 enum BazelTargetStoreError: Error, LocalizedError {
     case unknownBSPURI(URI)
     case unknownBazelLabel(String)
@@ -79,6 +85,9 @@ final class BazelTargetStoreImpl: BazelTargetStore {
         self.bazelTargetQuerier = bazelTargetQuerier
     }
 
+    /// Maps the list of supported platforms to the list of top-level labels of said platform.
+    var platformsToTopLevelLabelsMap: [String: [String]] = [:]
+
     /// Converts a BSP BuildTarget URI to its underlying Bazel target label.
     func bazelTargetLabel(forBSPURI uri: URI) throws -> String {
         guard let label = bspURIsToBazelLabelsMap[uri] else {
@@ -121,7 +130,7 @@ final class BazelTargetStoreImpl: BazelTargetStore {
 
     /// Provides the bazel label containing **platform information** for a given BSP URI.
     /// This is used to determine the correct set of compiler flags for the target / platform combo.
-    func platformBuildLabel(forBSPURI uri: URI) throws -> (String, TopLevelRuleType) {
+    func platformBuildLabelInfo(forBSPURI uri: URI) throws -> BazelTargetPlatformInfo {
         let bazelLabel = try bazelTargetLabel(forBSPURI: uri)
         let parents = try bazelLabelToParents(forBazelLabel: bazelLabel)
         // FIXME: When a target can compile to multiple platforms, the way Xcode handles it is by selecting
@@ -132,9 +141,9 @@ final class BazelTargetStoreImpl: BazelTargetStore {
         let baseSuffix = initializedConfig.baseConfig.buildTestSuffix
         let platformPlaceholder = initializedConfig.baseConfig.buildTestPlatformPlaceholder
         let platformBuildSuffix = baseSuffix.replacingOccurrences(of: platformPlaceholder, with: rule.platform)
-        return (
-            "\(bazelLabel)\(platformBuildSuffix)",
-            rule
+        return BazelTargetPlatformInfo(
+            buildTestLabel: "\(bazelLabel)\(platformBuildSuffix)",
+            parentRuleType: rule,
         )
     }
 
@@ -191,6 +200,7 @@ final class BazelTargetStoreImpl: BazelTargetStore {
         }
         for (target, ruleType) in topLevelTargetData {
             topLevelLabelToRuleMap[target] = ruleType
+            platformsToTopLevelLabelsMap[ruleType.platform, default: []].append(target)
         }
 
         // We need to now map which targets belong to which top-level apps,
@@ -254,6 +264,7 @@ final class BazelTargetStoreImpl: BazelTargetStore {
         bazelLabelToParentsMap = [:]
         availableBazelLabels = []
         topLevelLabelToRuleMap = [:]
+        platformsToTopLevelLabelsMap = [:]
         bazelTargetQuerier.clearCache()
         cachedTargets = nil
     }

Sources/SourceKitBazelBSP/RequestHandlers/DidInitializeHandler.swift

@@ -23,6 +23,10 @@ import LanguageServerProtocol
 
 private let logger = makeFileLevelBSPLogger()
 
+protocol DidInitializeObserver: AnyObject {
+    func didInitializeHandlerFinishedPreparations()
+}
+
 /// Handles the `build/initialized` notification.
 ///
 /// This is called right after returning from the `initialize` request.
@@ -31,12 +35,15 @@ final class DidInitializeHandler: @unchecked Sendable {
 
     private let initializedConfig: InitializedServerConfig
     private let commandRunner: CommandRunner
+    private let observers: [DidInitializeObserver]
 
     init(
         initializedConfig: InitializedServerConfig,
+        observers: [DidInitializeObserver],
         commandRunner: CommandRunner = ShellCommandRunner(),
     ) {
         self.initializedConfig = initializedConfig
+        self.observers = observers
         self.commandRunner = commandRunner
     }
 
@@ -58,6 +65,7 @@ final class DidInitializeHandler: @unchecked Sendable {
                 return
             }
             guard initializedConfig.aqueryOutputBase != initializedConfig.outputBase else {
+                self.notifyObservers()
                 return
             }
             // FIXME: We have to warm up the aqueries *after* the build, otherwise we can run
@@ -70,7 +78,12 @@ final class DidInitializeHandler: @unchecked Sendable {
             )
             aquery?.setTerminationHandler { code in
                 logger.info("Finished warming up the aquery output base! (status code: \(code))")
+                self.notifyObservers()
             }
         }
     }
+
+    func notifyObservers() {
+        observers.forEach { $0.didInitializeHandlerFinishedPreparations() }
+    }
 }

Sources/SourceKitBazelBSP/RequestHandlers/InitializeHandler.swift

@@ -112,9 +112,11 @@ final class InitializeHandler {
         // Collecting the rest of the env's details
         let devDir: String = try commandRunner.run("xcode-select --print-path")
         let toolchain = try getToolchainPath(with: commandRunner)
+        let sdkRootPaths: [String: String] = getSDKRootPaths(with: commandRunner)
 
         logger.debug("devDir: \(devDir)")
         logger.debug("toolchain: \(toolchain)")
+        logger.debug("sdkRootPaths: \(sdkRootPaths)")
 
         return InitializedServerConfig(
             baseConfig: baseConfig,
@@ -123,7 +125,8 @@ final class InitializeHandler {
             outputPath: outputPath,
             devDir: devDir,
             devToolchainPath: toolchain,
-            executionRoot: executionRoot
+            executionRoot: executionRoot,
+            sdkRootPaths: sdkRootPaths
         )
     }
 
@@ -140,6 +143,18 @@ final class InitializeHandler {
         return String(toolchain)
     }
 
+    func getSDKRootPaths(with commandRunner: CommandRunner) -> [String: String] {
+        let supportedSDKTypes = Set(TopLevelRuleType.allCases.map { $0.sdkName}).sorted()
+        let sdkRootPaths: [String: String] = supportedSDKTypes.reduce(into: [:]) { result, sdkType in
+            // This will fail if the user doesn't have the SDK installed, which is fine.
+            guard let sdkRootPath: String? = try? commandRunner.run("xcrun --sdk \(sdkType) --show-sdk-path") else {
+                return
+            }
+            result[sdkType] = sdkRootPath
+        }
+        return sdkRootPaths
+    }
+
     func buildResponse(
         fromRequest request: InitializeBuildRequest,
         and initializedConfig: InitializedServerConfig,

Sources/SourceKitBazelBSP/RequestHandlers/PrepareHandler.swift

@@ -69,7 +69,7 @@ final class PrepareHandler {
         do {
             let labels = try targetStore.stateLock.withLockUnchecked {
                 return try targetsToBuild.map {
-                    try targetStore.platformBuildLabel(forBSPURI: $0.uri).0
+                    try targetStore.platformBuildLabelInfo(forBSPURI: $0.uri).buildTestLabel
                 }
             }
             nonisolated(unsafe) let reply = reply

Sources/SourceKitBazelBSP/RequestHandlers/SKOptions/AqueryResult.swift

@@ -0,0 +1,58 @@
+// Copyright (c) 2025 Spotify AB.
+//
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+import BazelProtobufBindings
+import Foundation
+
+private let logger = makeFileLevelBSPLogger()
+
+enum AqueryResultError: Error, LocalizedError {
+    case duplicateTarget(label: String)
+    case duplicateAction(targetID: UInt32)
+
+    var errorDescription: String? {
+        switch self {
+        case .duplicateTarget(let label): return "Duplicate target found in the aquery! (\(label)) This can happen if a target gets different arguments depending on which top-level target builds it (on the same platform). Currently, the BSP expects the target to be stable in that sense."
+        case .duplicateAction(let targetID): return "Duplicate action ID found in the aquery! (\(targetID)) This is unexpected. Failing pre-emptively."
+        }
+    }
+}
+
+/// Small abstraction on top of Analysis_ActionGraphContainer to pre-aggregate the proto results.
+final class AqueryResult {
+    let targets: [String: Analysis_Target]
+    let actions: [UInt32: Analysis_Action]
+
+    init(results: Analysis_ActionGraphContainer) throws {
+        let targets: [String: Analysis_Target] = try results.targets.reduce(into: [:]) { result, target in
+            if result.keys.contains(target.label) {
+                throw AqueryResultError.duplicateTarget(label: target.label)
+            }
+            result[target.label] = target
+        }
+        let actions: [UInt32: Analysis_Action] = try results.actions.reduce(into: [:]) { result, action in
+            if result.keys.contains(action.targetID) {
+                throw AqueryResultError.duplicateAction(targetID: action.targetID)
+            }
+            result[action.targetID] = action
+        }
+        self.targets = targets
+        self.actions = actions
+    }
+}

Sources/SourceKitBazelBSP/RequestHandlers/SKOptions/BazelTargetAquerier.swift

@@ -37,29 +37,28 @@ enum BazelTargetAquerierError: Error, LocalizedError {
 final class BazelTargetAquerier {
 
     private let commandRunner: CommandRunner
-    private var queryCache = [String: Analysis_ActionGraphContainer]()
+    private var queryCache = [String: AqueryResult]()
 
     init(commandRunner: CommandRunner = ShellCommandRunner()) {
         self.commandRunner = commandRunner
     }
 
     func aquery(
-        target: String,
-        filteringFor: String,
+        targets: [String],
         config: InitializedServerConfig,
         mnemonics: Set<String>,
         additionalFlags: [String]
-    ) throws -> Analysis_ActionGraphContainer {
+    ) throws -> AqueryResult {
         guard !mnemonics.isEmpty else {
             throw BazelTargetAquerierError.noMnemonics
         }
 
         let mnemonicsFilter = mnemonics.sorted().joined(separator: "|")
-        let depsQuery = BazelTargetQuerier.queryDepsString(forTargets: [target])
+        let depsQuery = BazelTargetQuerier.queryDepsString(forTargets: targets)
 
         let otherFlags = additionalFlags.joined(separator: " ") + " --output proto"
-        let cmd = "aquery \"mnemonic('\(mnemonicsFilter)', filter(\(filteringFor), \(depsQuery)))\" \(otherFlags)"
-        logger.info("Processing aquery request for \(target), filtering for \(filteringFor)")
+        let cmd = "aquery \"mnemonic('\(mnemonicsFilter)', \(depsQuery))\" \(otherFlags)"
+        logger.info("Processing aquery request for \(targets)")
 
         if let cached = queryCache[cmd] {
             logger.debug("Returning cached results")
@@ -75,12 +74,13 @@ final class BazelTargetAquerier {
         )
 
         let parsedOutput = try BazelProtobufBindings.parseActionGraph(data: output)
+        let aqueryResult = try AqueryResult(results: parsedOutput)
 
         logger.debug("ActionGraphContainer parsed \(parsedOutput.actions.count) actions")
 
-        queryCache[cmd] = parsedOutput
+        queryCache[cmd] = aqueryResult
 
-        return parsedOutput
+        return aqueryResult
     }
 
     func clearCache() {