Refactor and cleanup of cquery processing

Test: USE_BAZEL_ANALYSIS=1 m libc
Change-Id: Iaf9a92e84d39c132e2444a8aaafd79505a12b8ec
diff --git a/android/Android.bp b/android/Android.bp
index f17a8a0..34fa811 100644
--- a/android/Android.bp
+++ b/android/Android.bp
@@ -12,6 +12,7 @@
         "soong",
         "soong-android-soongconfig",
         "soong-bazel",
+        "soong-cquery",
         "soong-shared",
         "soong-ui-metrics_proto",
     ],
diff --git a/android/bazel_handler.go b/android/bazel_handler.go
index 6675840..bbec389 100644
--- a/android/bazel_handler.go
+++ b/android/bazel_handler.go
@@ -26,6 +26,7 @@
 	"strings"
 	"sync"
 
+	"android/soong/bazel/cquery"
 	"github.com/google/blueprint/bootstrap"
 
 	"android/soong/bazel"
@@ -43,7 +44,7 @@
 // Map key to describe bazel cquery requests.
 type cqueryKey struct {
 	label       string
-	requestType CqueryRequestType
+	requestType cquery.RequestType
 	archType    ArchType
 }
 
@@ -53,14 +54,15 @@
 	// has been queued to be run later.
 
 	// Returns result files built by building the given bazel target label.
-	GetAllFiles(label string, archType ArchType) ([]string, bool)
+	GetOutputFiles(label string, archType ArchType) ([]string, bool)
 
 	// Returns object files produced by compiling the given cc-related target.
 	// Retrieves these files from Bazel's CcInfo provider.
 	GetCcObjectFiles(label string, archType ArchType) ([]string, bool)
 
-	// Returns the results of GetAllFiles and GetCcObjectFiles in a single query (in that order).
-	GetAllFilesAndCcObjectFiles(label string, archType ArchType) ([]string, []string, bool)
+	// TODO(cparsons): Other cquery-related methods should be added here.
+	// Returns the results of GetOutputFiles and GetCcObjectFiles in a single query (in that order).
+	GetOutputFilesAndCcObjectFiles(label string, archType ArchType) ([]string, []string, bool)
 
 	// ** End cquery methods
 
@@ -109,7 +111,7 @@
 	AllFiles map[string][]string
 }
 
-func (m MockBazelContext) GetAllFiles(label string, archType ArchType) ([]string, bool) {
+func (m MockBazelContext) GetOutputFiles(label string, archType ArchType) ([]string, bool) {
 	result, ok := m.AllFiles[label]
 	return result, ok
 }
@@ -119,7 +121,7 @@
 	return result, ok
 }
 
-func (m MockBazelContext) GetAllFilesAndCcObjectFiles(label string, archType ArchType) ([]string, []string, bool) {
+func (m MockBazelContext) GetOutputFilesAndCcObjectFiles(label string, archType ArchType) ([]string, []string, bool) {
 	result, ok := m.AllFiles[label]
 	return result, result, ok
 }
@@ -142,43 +144,42 @@
 
 var _ BazelContext = MockBazelContext{}
 
-func (bazelCtx *bazelContext) GetAllFiles(label string, archType ArchType) ([]string, bool) {
-	result, ok := bazelCtx.cquery(label, getAllFiles, archType)
+func (bazelCtx *bazelContext) GetOutputFiles(label string, archType ArchType) ([]string, bool) {
+	rawString, ok := bazelCtx.cquery(label, cquery.GetOutputFiles, archType)
+	var ret []string
 	if ok {
-		bazelOutput := strings.TrimSpace(result)
-		return strings.Split(bazelOutput, ", "), true
-	} else {
-		return nil, false
+		bazelOutput := strings.TrimSpace(rawString)
+		ret = cquery.GetOutputFiles.ParseResult(bazelOutput).([]string)
 	}
+	return ret, ok
 }
 
 func (bazelCtx *bazelContext) GetCcObjectFiles(label string, archType ArchType) ([]string, bool) {
-	result, ok := bazelCtx.cquery(label, getCcObjectFiles, archType)
+	rawString, ok := bazelCtx.cquery(label, cquery.GetCcObjectFiles, archType)
+	var returnResult []string
 	if ok {
-		bazelOutput := strings.TrimSpace(result)
-		return strings.Split(bazelOutput, ", "), true
-	} else {
-		return nil, false
+		bazelOutput := strings.TrimSpace(rawString)
+		returnResult = cquery.GetCcObjectFiles.ParseResult(bazelOutput).([]string)
 	}
+	return returnResult, ok
 }
 
-func (bazelCtx *bazelContext) GetAllFilesAndCcObjectFiles(label string, archType ArchType) ([]string, []string, bool) {
-	var allFiles []string
+func (bazelCtx *bazelContext) GetOutputFilesAndCcObjectFiles(label string, archType ArchType) ([]string, []string, bool) {
+	var outputFiles []string
 	var ccObjects []string
 
-	result, ok := bazelCtx.cquery(label, getAllFilesAndCcObjectFiles, archType)
+	result, ok := bazelCtx.cquery(label, cquery.GetOutputFilesAndCcObjectFiles, archType)
 	if ok {
 		bazelOutput := strings.TrimSpace(result)
-		splitString := strings.Split(bazelOutput, "|")
-		allFilesString := splitString[0]
-		ccObjectsString := splitString[1]
-		allFiles = strings.Split(allFilesString, ", ")
-		ccObjects = strings.Split(ccObjectsString, ", ")
+		returnResult := cquery.GetOutputFilesAndCcObjectFiles.ParseResult(bazelOutput).(cquery.GetOutputFilesAndCcObjectFiles_Result)
+		outputFiles = returnResult.OutputFiles
+		ccObjects = returnResult.CcObjectFiles
 	}
-	return allFiles, ccObjects, ok
+
+	return outputFiles, ccObjects, ok
 }
 
-func (n noopBazelContext) GetAllFiles(label string, archType ArchType) ([]string, bool) {
+func (n noopBazelContext) GetOutputFiles(label string, archType ArchType) ([]string, bool) {
 	panic("unimplemented")
 }
 
@@ -186,7 +187,7 @@
 	panic("unimplemented")
 }
 
-func (n noopBazelContext) GetAllFilesAndCcObjectFiles(label string, archType ArchType) ([]string, []string, bool) {
+func (n noopBazelContext) GetOutputFilesAndCcObjectFiles(label string, archType ArchType) ([]string, []string, bool) {
 	panic("unimplemented")
 }
 
@@ -260,7 +261,7 @@
 // If the given request was already made (and the results are available), then
 // returns (result, true). If the request is queued but no results are available,
 // then returns ("", false).
-func (context *bazelContext) cquery(label string, requestType CqueryRequestType,
+func (context *bazelContext) cquery(label string, requestType cquery.RequestType,
 	archType ArchType) (string, bool) {
 	key := cqueryKey{label, requestType, archType}
 	if result, ok := context.results[key]; ok {
@@ -485,38 +486,66 @@
 		strings.Join(deps_arm, ",\n            ")))
 }
 
+func indent(original string) string {
+	result := ""
+	for _, line := range strings.Split(original, "\n") {
+		result += "  " + line + "\n"
+	}
+	return result
+}
+
 // Returns the file contents of the buildroot.cquery file that should be used for the cquery
 // expression in order to obtain information about buildroot and its dependencies.
 // The contents of this file depend on the bazelContext's requests; requests are enumerated
 // and grouped by their request type. The data retrieved for each label depends on its
 // request type.
 func (context *bazelContext) cqueryStarlarkFileContents() []byte {
+	requestTypeToCqueryIdEntries := map[cquery.RequestType][]string{}
+	for val, _ := range context.requests {
+		cqueryId := getCqueryId(val)
+		mapEntryString := fmt.Sprintf("%q : True", cqueryId)
+		requestTypeToCqueryIdEntries[val.requestType] =
+			append(requestTypeToCqueryIdEntries[val.requestType], mapEntryString)
+	}
+	labelRegistrationMapSection := ""
+	functionDefSection := ""
+	mainSwitchSection := ""
+
+	mapDeclarationFormatString := `
+%s = {
+  %s
+}
+`
+	functionDefFormatString := `
+def %s(target):
+%s
+`
+	mainSwitchSectionFormatString := `
+  if id_string in %s:
+    return id_string + ">>" + %s(target)
+`
+
+	for _, requestType := range cquery.RequestTypes {
+		labelMapName := requestType.Name() + "_Labels"
+		functionName := requestType.Name() + "_Fn"
+		labelRegistrationMapSection += fmt.Sprintf(mapDeclarationFormatString,
+			labelMapName,
+			strings.Join(requestTypeToCqueryIdEntries[requestType], ",\n  "))
+		functionDefSection += fmt.Sprintf(functionDefFormatString,
+			functionName,
+			indent(requestType.StarlarkFunctionBody()))
+		mainSwitchSection += fmt.Sprintf(mainSwitchSectionFormatString,
+			labelMapName, functionName)
+	}
+
 	formatString := `
 # This file is generated by soong_build. Do not edit.
-getAllFilesLabels = {
-  %s
-}
 
-getCcObjectFilesLabels = {
-  %s
-}
+# Label Map Section
+%s
 
-getAllFilesAndCcObjectFilesLabels = {
-  %s
-}
-
-def get_all_files(target):
-  return [f.path for f in target.files.to_list()]
-
-def get_cc_object_files(target):
-  result = []
-  linker_inputs = providers(target)["CcInfo"].linking_context.linker_inputs.to_list()
-
-  for linker_input in linker_inputs:
-    for library in linker_input.libraries:
-      for object in library.objects:
-        result += [object.path]
-  return result
+# Function Def Section
+%s
 
 def get_arch(target):
   buildoptions = build_options(target)
@@ -536,39 +565,16 @@
 
 def format(target):
   id_string = str(target.label) + "|" + get_arch(target)
-  if id_string in getAllFilesLabels:
-    return id_string + ">>" + ', '.join(get_all_files(target))
-  elif id_string in getCcObjectFilesLabels:
-    return id_string + ">>" + ', '.join(get_cc_object_files(target))
-  elif id_string in getAllFilesAndCcObjectFilesLabels:
-    return id_string + ">>" + ', '.join(get_all_files(target)) + "|" + ', '.join(get_cc_object_files(target))
-  else:
-    # This target was not requested via cquery, and thus must be a dependency
-    # of a requested target.
-    return id_string + ">>NONE"
+
+  # Main switch section
+  %s
+  # This target was not requested via cquery, and thus must be a dependency
+  # of a requested target.
+  return id_string + ">>NONE"
 `
-	var getAllFilesDeps []string = nil
-	var getCcObjectFilesDeps []string = nil
-	var getAllFilesAndCcObjectFilesDeps []string = nil
 
-	for val, _ := range context.requests {
-		labelWithArch := getCqueryId(val)
-		mapEntryString := fmt.Sprintf("%q : True", labelWithArch)
-		switch val.requestType {
-		case getAllFiles:
-			getAllFilesDeps = append(getAllFilesDeps, mapEntryString)
-		case getCcObjectFiles:
-			getCcObjectFilesDeps = append(getCcObjectFilesDeps, mapEntryString)
-		case getAllFilesAndCcObjectFiles:
-			getAllFilesAndCcObjectFilesDeps = append(getAllFilesAndCcObjectFilesDeps, mapEntryString)
-		}
-	}
-	getAllFilesDepsString := strings.Join(getAllFilesDeps, ",\n  ")
-	getCcObjectFilesDepsString := strings.Join(getCcObjectFilesDeps, ",\n  ")
-	getAllFilesAndCcObjectFilesDepsString := strings.Join(getAllFilesAndCcObjectFilesDeps, ",\n  ")
-
-	return []byte(fmt.Sprintf(formatString, getAllFilesDepsString, getCcObjectFilesDepsString,
-		getAllFilesAndCcObjectFilesDepsString))
+	return []byte(fmt.Sprintf(formatString, labelRegistrationMapSection, functionDefSection,
+		mainSwitchSection))
 }
 
 // Returns a workspace-relative path containing build-related metadata required
diff --git a/bazel/cquery/Android.bp b/bazel/cquery/Android.bp
new file mode 100644
index 0000000..3a71e9c
--- /dev/null
+++ b/bazel/cquery/Android.bp
@@ -0,0 +1,14 @@
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+bootstrap_go_package {
+    name: "soong-cquery",
+    pkgPath: "android/soong/bazel/cquery",
+    srcs: [
+        "request_type.go",
+    ],
+    pluginFor: [
+        "soong_build",
+    ],
+}
diff --git a/bazel/cquery/request_type.go b/bazel/cquery/request_type.go
new file mode 100644
index 0000000..864db3d
--- /dev/null
+++ b/bazel/cquery/request_type.go
@@ -0,0 +1,110 @@
+package cquery
+
+import (
+	"strings"
+)
+
+var (
+	GetOutputFiles                 RequestType = &getOutputFilesRequestType{}
+	GetCcObjectFiles               RequestType = &getCcObjectFilesRequestType{}
+	GetOutputFilesAndCcObjectFiles RequestType = &getOutputFilesAndCcObjectFilesType{}
+)
+
+type GetOutputFilesAndCcObjectFiles_Result struct {
+	OutputFiles   []string
+	CcObjectFiles []string
+}
+
+var RequestTypes []RequestType = []RequestType{
+	GetOutputFiles, GetCcObjectFiles, GetOutputFilesAndCcObjectFiles}
+
+type RequestType interface {
+	// Name returns a string name for this request type. Such request type names must be unique,
+	// and must only consist of alphanumeric characters.
+	Name() string
+
+	// StarlarkFunctionBody returns a straark function body to process this request type.
+	// The returned string is the body of a Starlark function which obtains
+	// all request-relevant information about a target and returns a string containing
+	// this information.
+	// The function should have the following properties:
+	//   - `target` is the only parameter to this function (a configured target).
+	//   - The return value must be a string.
+	//   - The function body should not be indented outside of its own scope.
+	StarlarkFunctionBody() string
+
+	// ParseResult returns a value obtained by parsing the result of the request's Starlark function.
+	// The given rawString must correspond to the string output which was created by evaluating the
+	// Starlark given in StarlarkFunctionBody.
+	// The type of this value depends on the request type; it is up to the caller to
+	// cast to the correct type.
+	ParseResult(rawString string) interface{}
+}
+
+type getOutputFilesRequestType struct{}
+
+func (g getOutputFilesRequestType) Name() string {
+	return "getOutputFiles"
+}
+
+func (g getOutputFilesRequestType) StarlarkFunctionBody() string {
+	return "return ', '.join([f.path for f in target.files.to_list()])"
+}
+
+func (g getOutputFilesRequestType) ParseResult(rawString string) interface{} {
+	return strings.Split(rawString, ", ")
+}
+
+type getCcObjectFilesRequestType struct{}
+
+func (g getCcObjectFilesRequestType) Name() string {
+	return "getCcObjectFiles"
+}
+
+func (g getCcObjectFilesRequestType) StarlarkFunctionBody() string {
+	return `
+result = []
+linker_inputs = providers(target)["CcInfo"].linking_context.linker_inputs.to_list()
+
+for linker_input in linker_inputs:
+  for library in linker_input.libraries:
+    for object in library.objects:
+      result += [object.path]
+return ', '.join(result)`
+}
+
+func (g getCcObjectFilesRequestType) ParseResult(rawString string) interface{} {
+	return strings.Split(rawString, ", ")
+}
+
+type getOutputFilesAndCcObjectFilesType struct{}
+
+func (g getOutputFilesAndCcObjectFilesType) Name() string {
+	return "getOutputFilesAndCcObjectFiles"
+}
+
+func (g getOutputFilesAndCcObjectFilesType) StarlarkFunctionBody() string {
+	return `
+outputFiles = [f.path for f in target.files.to_list()]
+
+ccObjectFiles = []
+linker_inputs = providers(target)["CcInfo"].linking_context.linker_inputs.to_list()
+
+for linker_input in linker_inputs:
+  for library in linker_input.libraries:
+    for object in library.objects:
+      ccObjectFiles += [object.path]
+return ', '.join(outputFiles) + "|" + ', '.join(ccObjectFiles)`
+}
+
+func (g getOutputFilesAndCcObjectFilesType) ParseResult(rawString string) interface{} {
+	var outputFiles []string
+	var ccObjects []string
+
+	splitString := strings.Split(rawString, "|")
+	outputFilesString := splitString[0]
+	ccObjectsString := splitString[1]
+	outputFiles = strings.Split(outputFilesString, ", ")
+	ccObjects = strings.Split(ccObjectsString, ", ")
+	return GetOutputFilesAndCcObjectFiles_Result{outputFiles, ccObjects}
+}
diff --git a/cc/library.go b/cc/library.go
index 6a3b876..22a36c6 100644
--- a/cc/library.go
+++ b/cc/library.go
@@ -415,38 +415,39 @@
 
 func (handler *staticLibraryBazelHandler) generateBazelBuildActions(ctx android.ModuleContext, label string) bool {
 	bazelCtx := ctx.Config().BazelContext
-	outputPaths, objPaths, ok := bazelCtx.GetAllFilesAndCcObjectFiles(label, ctx.Arch().ArchType)
-	if ok {
-		if len(outputPaths) != 1 {
-			// TODO(cparsons): This is actually expected behavior for static libraries with no srcs.
-			// We should support this.
-			ctx.ModuleErrorf("expected exactly one output file for '%s', but got %s", label, objPaths)
-			return false
-		}
-		outputFilePath := android.PathForBazelOut(ctx, outputPaths[0])
-		handler.module.outputFile = android.OptionalPathForPath(outputFilePath)
-
-		objFiles := make(android.Paths, len(objPaths))
-		for i, objPath := range objPaths {
-			objFiles[i] = android.PathForBazelOut(ctx, objPath)
-		}
-		objects := Objects{
-			objFiles: objFiles,
-		}
-
-		ctx.SetProvider(StaticLibraryInfoProvider, StaticLibraryInfo{
-			StaticLibrary: outputFilePath,
-			ReuseObjects:  objects,
-			Objects:       objects,
-
-			// TODO(cparsons): Include transitive static libraries in this provider to support
-			// static libraries with deps.
-			TransitiveStaticLibrariesForOrdering: android.NewDepSetBuilder(android.TOPOLOGICAL).
-				Direct(outputFilePath).
-				Build(),
-		})
-		handler.module.outputFile = android.OptionalPathForPath(android.PathForBazelOut(ctx, objPaths[0]))
+	outputPaths, objPaths, ok := bazelCtx.GetOutputFilesAndCcObjectFiles(label, ctx.Arch().ArchType)
+	if !ok {
+		return ok
 	}
+	if len(outputPaths) != 1 {
+		// TODO(cparsons): This is actually expected behavior for static libraries with no srcs.
+		// We should support this.
+		ctx.ModuleErrorf("expected exactly one output file for '%s', but got %s", label, objPaths)
+		return false
+	}
+	outputFilePath := android.PathForBazelOut(ctx, outputPaths[0])
+	handler.module.outputFile = android.OptionalPathForPath(outputFilePath)
+
+	objFiles := make(android.Paths, len(objPaths))
+	for i, objPath := range objPaths {
+		objFiles[i] = android.PathForBazelOut(ctx, objPath)
+	}
+	objects := Objects{
+		objFiles: objFiles,
+	}
+
+	ctx.SetProvider(StaticLibraryInfoProvider, StaticLibraryInfo{
+		StaticLibrary: outputFilePath,
+		ReuseObjects:  objects,
+		Objects:       objects,
+
+		// TODO(cparsons): Include transitive static libraries in this provider to support
+		// static libraries with deps.
+		TransitiveStaticLibrariesForOrdering: android.NewDepSetBuilder(android.TOPOLOGICAL).
+			Direct(outputFilePath).
+			Build(),
+	})
+	handler.module.outputFile = android.OptionalPathForPath(android.PathForBazelOut(ctx, objPaths[0]))
 	return ok
 }
 
diff --git a/genrule/genrule.go b/genrule/genrule.go
index 5349906..d9ec902 100644
--- a/genrule/genrule.go
+++ b/genrule/genrule.go
@@ -227,7 +227,7 @@
 // Returns true if information was available from Bazel, false if bazel invocation still needs to occur.
 func (c *Module) generateBazelBuildActions(ctx android.ModuleContext, label string) bool {
 	bazelCtx := ctx.Config().BazelContext
-	filePaths, ok := bazelCtx.GetAllFiles(label, ctx.Arch().ArchType)
+	filePaths, ok := bazelCtx.GetOutputFiles(label, ctx.Arch().ArchType)
 	if ok {
 		var bazelOutputFiles android.Paths
 		for _, bazelOutputFile := range filePaths {