Migrate action query from text to protobuf parsing

Package.swift

@@ -51,6 +51,8 @@ let package = Package(
             name: "SourceKitBazelBSPTests",
             dependencies: ["SourceKitBazelBSP"],
             resources: [
+                .copy("Resources/aquery.pb"),
+                .copy("Resources/aquery_objc.pb"),
                 .copy("Resources/streamdeps.pb"),
             ],
         ),

Sources/SourceKitBazelBSP/RequestHandlers/SKOptions/BazelTargetAquerier.swift

@@ -17,6 +17,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+import BazelProtobufBindings
 import Foundation
 
 private let logger = makeFileLevelBSPLogger()
@@ -33,12 +34,10 @@ enum BazelTargetAquerierError: Error, LocalizedError {
 
 /// Small abstraction to handle and cache the results of bazel _action queries_.
 /// FIXME: This is separate from BazelTargetQuerier because of the different output types, but we can unify these.
-///
-/// FIXME: Currently uses text outputs, should use proto instead so that we can organize and test this properly.
 final class BazelTargetAquerier {
 
     private let commandRunner: CommandRunner
-    private var queryCache = [String: String]()
+    private var queryCache = [String: Analysis_ActionGraphContainer]()
 
     init(commandRunner: CommandRunner = ShellCommandRunner()) {
         self.commandRunner = commandRunner
@@ -50,15 +49,15 @@ final class BazelTargetAquerier {
         config: InitializedServerConfig,
         mnemonics: Set<String>,
         additionalFlags: [String]
-    ) throws -> String {
+    ) throws -> Analysis_ActionGraphContainer {
         guard !mnemonics.isEmpty else {
             throw BazelTargetAquerierError.noMnemonics
         }
 
         let mnemonicsFilter = mnemonics.sorted().joined(separator: "|")
         let depsQuery = BazelTargetQuerier.queryDepsString(forTargets: [target])
 
-        let otherFlags = additionalFlags.joined(separator: " ")
+        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)")
 
@@ -68,11 +67,15 @@ final class BazelTargetAquerier {
         }
 
         // Run the aquery on the special index output base since that's where we will build at.
-        let output = try commandRunner.bazelIndexAction(initializedConfig: config, cmd: cmd)
+        let output = try commandRunner.bazelActionQuery(initializedConfig: config, cmd: cmd)
 
-        queryCache[cmd] = output
+        let parsedOutput = try BazelProtobufBindings.parseActionGraph(data: output)
 
-        return output
+        logger.debug("actionGraphContainer contains \(parsedOutput.actions.count, privacy: .public) actions.")
+
+        queryCache[cmd] = parsedOutput
+
+        return parsedOutput
     }
 
     func clearCache() {

Sources/SourceKitBazelBSP/RequestHandlers/SKOptions/CompilerArgumentsProcessor.swift

@@ -17,112 +17,46 @@
 // specific language governing permissions and limitations
 // under the License.
 
+import BazelProtobufBindings
 import Foundation
 import LanguageServerProtocol
 
 private let logger = makeFileLevelBSPLogger()
 
 enum CompilerArgumentsProcessor {
     // Parses and processes the compilation step for a given target from a larger aquery output.
-    // The parsing step is only necessary because BazelTargetAquerier operates on text. Should become unnecessary once we move to proto.
     static func extractAndProcessCompilerArgs(
-        fromAquery aqueryOutput: String,
+        fromAquery aqueryOutput: Analysis_ActionGraphContainer,
         bazelTarget: String,
         contentToQuery: String,
         language: Language,
         initializedConfig: InitializedServerConfig
     ) -> [String]? {
-        var lines: [String] = aqueryOutput.components(separatedBy: "\n")
-        var idx = -1
-        for (i, line) in lines.enumerated() {
-            var actionPrefix = "action 'Compiling "
-            if language == .swift {
-                actionPrefix += "Swift module \(bazelTarget)'"
-            } else {
-                actionPrefix += "\(contentToQuery)'"
-            }
-            if line.hasPrefix(actionPrefix) {
-                if language == .swift {
-                    if lines[i + 1] != "  Mnemonic: SwiftCompile" { continue }
-                } else {
-                    if lines[i + 1] != "  Mnemonic: ObjcCompile" { continue }
-                    if lines[i + 2] != "  Target: \(bazelTarget)" { continue }
-                }
-                idx = i
-                break
-            }
-        }
-        if idx == -1 {
-            logger.error("No module entry found in the aquery for \(contentToQuery)")
+        guard let target = aqueryOutput.targets.first(where: { $0.label == bazelTarget }) else {
+            logger.debug("\(CompilerArgError.targetNotFound(bazelTarget)), privacy: .public")
             return nil
         }
-        // now, get the first index of the line that starts with "Command Line: ("
-        lines = Array(lines.dropFirst(idx + 1))
-        idx = -1
-        for (i, line) in lines.enumerated() {
-            if line.starts(with: "  Command Line: (") {
-                idx = i
-                // Also skip the swiftc line for swift targets
-                if lines[idx + 1].contains("swiftc") { idx += 1 }
-                break
-            }
-        }
-        if idx == -1 {
-            logger.error("No command line entry found")
+        guard let action = aqueryOutput.actions.first(where: { $0.targetID == target.id }) else {
+            logger.debug("\(CompilerArgError.actionNotFound(bazelTarget)), privacy: .public")
             return nil
         }
-        logger.info("Found command line entry at \(idx)")
-
-        // now, find where the arguments end
-        lines = Array(lines.dropFirst(idx + 1))
-
-        idx = -1
-        for (i, line) in lines.enumerated() {
-            if line.starts(with: "#") {
-                idx = i
-                break
-            }
-        }
-        if idx == -1 {
-            logger.error("Couldn't find where the args end")
-            return nil
-        }
-
-        logger.info("Found where the args end at \(idx)")
-        lines = Array(lines.dropLast(lines.count - idx - 1))
-
-        // the spaced lines are the compiler arguments
-        lines = lines.filter { $0.starts(with: "    ") }
 
-        lines = lines.map { $0.trimmingCharacters(in: .whitespacesAndNewlines) }
-        lines = lines.map { $0.hasSuffix(" \\") ? String($0.dropLast(2)) : $0 }
-        // remove the trailing ) from the last line
-        lines[lines.count - 1] = String(lines[lines.count - 1].dropLast())
-
-        // some args are wrapped in single quotes for some reason
-        for i in 0..<lines.count {
-            if lines[i].hasPrefix("'"), lines[i].hasSuffix("'") {
-                lines[i] = String(lines[i].dropFirst().dropLast())
-            }
-        }
+        let rawArguments = action.arguments
 
         let processedArgs = _processCompilerArguments(
-            rawArguments: lines,
+            rawArguments: rawArguments,
             contentToQuery: contentToQuery,
             language: language,
             initializedConfig: initializedConfig
         )
 
-        lines = processedArgs
-
         logger.info("Finished processing compiler arguments")
         logger.logFullObjectInMultipleLogMessages(
             level: .debug,
             header: "Parsed compiler arguments",
-            lines.joined(separator: "\n"),
+            processedArgs.joined(separator: "\n"),
         )
-
-        return lines
+        return processedArgs
     }
 
     /// Processes compiler arguments for the LSP by removing unnecessary arguments and replacing placeholders.
@@ -154,6 +88,15 @@ enum CompilerArgumentsProcessor {
         while index < count {
             let arg = rawArguments[index]
 
+            // Skip worker and swiftc
+            if arg.hasSuffix("/worker"),
+                index + 1 < count,
+                rawArguments[index + 1] == "swiftc"
+            {
+                index += 2
+                continue
+            }
+
             // Skip wrapped arguments. These don't work for some reason
             if arg.hasPrefix("-Xwrapped-swift") {
                 index += 1
@@ -266,3 +209,14 @@ enum CompilerArgumentsProcessor {
         lines[idx + 1] = new
     }
 }
+
+enum CompilerArgError: Error, CustomStringConvertible {
+    case targetNotFound(String)
+    case actionNotFound(String)
+    var description: String {
+        switch self {
+        case .targetNotFound(let target): "Target: \(target) not found."
+        case .actionNotFound(let target): "Action not found for target: \(target)."
+        }
+    }
+}

Sources/SourceKitBazelBSP/SharedUtils/Shell/CommandRunner.swift

@@ -52,6 +52,18 @@ extension CommandRunner {
         )
     }
 
+    /// A  bazel action query command, but at this BSP's special output base and taking into account the special index flags.
+    func bazelActionQuery(initializedConfig: InitializedServerConfig, cmd: String) throws -> Data {
+        let baseCommand = [
+            initializedConfig.baseConfig.bazelWrapper,
+            "--output_base=\(initializedConfig.outputBase)",
+            cmd,
+        ]
+        let additionalFlags = initializedConfig.baseConfig.indexFlags
+        let command = (baseCommand + additionalFlags).joined(separator: " ")
+        return try run(command, cwd: initializedConfig.rootUri)
+    }
+
     /// A regular bazel command, but at this BSP's special output base and taking into account the special index flags.
     func bazelIndexAction(
         baseConfig: BaseServerConfig,