Migrate action query from text to protobuf parsing (#40)

* replacing plain text parsing to proto parsing for action query results

  - append command string with "--output proto"
  - added CompilerArgError for target not found when processing compiler arguments
  - added back skip `swiftc` and `worker` arguments

* updated tests for action query

  - added aquery.pb and aquery_objc.pb as mocking data
  - updated expected aquery by append "--output proto"

* rename the command for action query

* fix the bug - the clang compiler call is included in the arguments list

* use opaque type for Data and String return type

* removed the error and using simple strings, improve performance with setting index only

* fix test

* change log to private

* use generic instead of opaque type, then call static func convert

Author: sean7218

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/BuildTargets/BazelTargetQuerier.swift

@@ -123,7 +123,7 @@ final class BazelTargetQuerier {
 
         // We run this one on the main output base since it's not related to the actual indexing bits
         let cmd = "query \"kind('\(kindsFilter)', \(depsQuery))\" --output streamed_proto"
-        let output = try commandRunner.run(config.bazelWrapper + " " + cmd, cwd: rootUri)
+        let output: Data = try commandRunner.run(config.bazelWrapper + " " + cmd, cwd: rootUri)
 
         logger.debug("Finished querying, building result Protobuf")
 

Sources/SourceKitBazelBSP/RequestHandlers/InitializeHandler.swift

@@ -92,7 +92,7 @@ final class InitializeHandler {
         logger.debug("outputBase: \(outputBase)")
 
         // Now, get the full output path based on the above output base.
-        let outputPath = try commandRunner.bazelIndexAction(
+        let outputPath: String = try commandRunner.bazelIndexAction(
             baseConfig: baseConfig,
             outputBase: outputBase,
             cmd: "info output_path",

Sources/SourceKitBazelBSP/RequestHandlers/PrepareHandler.swift

@@ -79,7 +79,7 @@ final class PrepareHandler {
         logger.info("Will build \(labelsToBuild.joined(separator: ", "))")
 
         // Build the provided targets, on our special output base and taking into account special index flags.
-        _ = try commandRunner.bazelIndexAction(
+        let _: String = try commandRunner.bazelIndexAction(
             initializedConfig: initializedConfig,
             cmd: "build \(labelsToBuild.joined(separator: " "))"
         )

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: Data = try commandRunner.bazelIndexAction(initializedConfig: config, cmd: cmd)
 
-        queryCache[cmd] = output
+        let parsedOutput = try BazelProtobufBindings.parseActionGraph(data: output)
 
-        return output
+        logger.debug("actionGraphContainer count \(parsedOutput.actions.count, privacy: .private)")
+
+        queryCache[cmd] = parsedOutput
+
+        return parsedOutput
     }
 
     func clearCache() {

Sources/SourceKitBazelBSP/RequestHandlers/SKOptions/CompilerArgumentsProcessor.swift

@@ -17,114 +17,48 @@
 // 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,
         sdkRoot: String,
         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)")
-            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 target = aqueryOutput.targets.first(where: { $0.label == bazelTarget }) else {
+            logger.debug("Target: \(bazelTarget, privacy: .private) not found.")
             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")
+        guard let action = aqueryOutput.actions.first(where: { $0.targetID == target.id }) else {
+            logger.debug("Action for \(bazelTarget, privacy: .private) not found.")
             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,
             sdkRoot: sdkRoot,
             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.
@@ -153,6 +87,13 @@ enum CompilerArgumentsProcessor {
         var index = 0
         let count = rawArguments.count
 
+        // drop clang and swiftc compiler call
+        switch language {
+        case .swift: index = 2
+        case .objective_c: index = 1
+        default: break
+        }
+
         while index < count {
             let arg = rawArguments[index]
 

Sources/SourceKitBazelBSP/SharedUtils/Shell/CommandRunner.swift

@@ -20,16 +20,10 @@
 import Foundation
 
 public protocol CommandRunner {
-    func run(_ cmd: String, cwd: String?) throws -> Data
+    func run<T: DataConvertible>(_ cmd: String, cwd: String?) throws -> T
 }
 
 extension CommandRunner {
-    func run(_ cmd: String, cwd: String?) throws -> String {
-        let data: Data = try run(cmd, cwd: cwd)
-        let str = String(data: data, encoding: .utf8) ?? ""
-        return str.trimmingCharacters(in: .whitespacesAndNewlines)
-    }
-
     func run(_ cmd: String) throws -> String {
         try run(cmd, cwd: nil)
     }
@@ -38,12 +32,12 @@ extension CommandRunner {
 // MARK: Bazel-related helpers
 
 extension CommandRunner {
-    func bazel(baseConfig: BaseServerConfig, rootUri: String, cmd: String) throws -> String {
+    func bazel<T: DataConvertible>(baseConfig: BaseServerConfig, rootUri: String, cmd: String) throws -> T {
         try run(baseConfig.bazelWrapper + " " + cmd, cwd: rootUri)
     }
 
     /// A regular bazel command, but at this BSP's special output base and taking into account the special index flags.
-    func bazelIndexAction(initializedConfig: InitializedServerConfig, cmd: String) throws -> String {
+    func bazelIndexAction<T: DataConvertible>(initializedConfig: InitializedServerConfig, cmd: String) throws -> T {
         return try bazelIndexAction(
             baseConfig: initializedConfig.baseConfig,
             outputBase: initializedConfig.outputBase,
@@ -53,12 +47,12 @@ extension CommandRunner {
     }
 
     /// A regular bazel command, but at this BSP's special output base and taking into account the special index flags.
-    func bazelIndexAction(
+    func bazelIndexAction<T: DataConvertible>(
         baseConfig: BaseServerConfig,
         outputBase: String,
         cmd: String,
         rootUri: String,
-    ) throws -> String {
+    ) throws -> T {
         let indexFlags = baseConfig.indexFlags
         let additionalFlags: String
         if indexFlags.isEmpty {
@@ -70,3 +64,20 @@ extension CommandRunner {
         return try bazel(baseConfig: baseConfig, rootUri: rootUri, cmd: cmd)
     }
 }
+
+public protocol DataConvertible {
+    static func convert(from data: Data) -> Self
+}
+
+extension String: DataConvertible {
+    public static func convert(from data: Data) -> Self {
+        let str = String(data: data, encoding: .utf8) ?? ""
+        return str.trimmingCharacters(in: .whitespacesAndNewlines)
+    }
+}
+
+extension Data: DataConvertible {
+    public static func convert(from data: Data) -> Self {
+        return data
+    }
+}

Sources/SourceKitBazelBSP/SharedUtils/Shell/ShellCommandRunner.swift

@@ -32,7 +32,7 @@ enum ShellCommandRunnerError: LocalizedError {
 }
 
 struct ShellCommandRunner: CommandRunner {
-    func run(_ cmd: String, cwd: String?) throws -> Data {
+    func run<T: DataConvertible>(_ cmd: String, cwd: String?) throws -> T {
         let task = Process()
         let stdout = Pipe()
         let stderr = Pipe()
@@ -64,6 +64,6 @@ struct ShellCommandRunner: CommandRunner {
             throw ShellCommandRunnerError.failed(cmd, stderrString)
         }
 
-        return data
+        return T.convert(from: data)
     }
 }