Refactor mixed builds to only take one pass

This large refactoring has both immense performance implications and
improves mixed builds complexity / usability. Summary:

1. Queueing calls to Bazel is done in a new mutator instead of a full
   soong_build pass. Normal soong_build flow is interrupted (via a
   functional hook in blueprint) to invoke bazel and parse its response.
2. Implementing mixed build support for additional modules is as simple
   as implementing MixedBuildsBuildable. In this interface, define the
   request that must be queued to Bazel, and then subsequently define
   how to handle the returned bazel cquery metadata.
3. Mixed builds consists of only a single pass. This greatly
   improves mixed build performance.

Result:
  A 33% runtime improvement on soong analysis phase with mixed builds.

Caveats:
  C++ BazelHandler handling still remains a bit of a mess; I did what
  I could within this CL's scope, but this may require additional cleanup.

Test: Treehugger
Test: Verified that aosp_arm ninja file is bit-for-bit identical with or
without this change.

Change-Id: I412d9c94d429105f4ebfafc84100d546069e6621
diff --git a/android/bazel.go b/android/bazel.go
index 67002ec..378eb6f 100644
--- a/android/bazel.go
+++ b/android/bazel.go
@@ -115,6 +115,27 @@
 	SetBaseModuleType(baseModuleType string)
 }
 
+// MixedBuildBuildable is an interface that module types should implement in order
+// to be "handled by Bazel" in a mixed build.
+type MixedBuildBuildable interface {
+	// IsMixedBuildSupported returns true if and only if this module should be
+	// "handled by Bazel" in a mixed build.
+	// This "escape hatch" allows modules with corner-case scenarios to opt out
+	// of being built with Bazel.
+	IsMixedBuildSupported(ctx BaseModuleContext) bool
+
+	// QueueBazelCall invokes request-queueing functions on the BazelContext
+	// so that these requests are handled when Bazel's cquery is invoked.
+	QueueBazelCall(ctx BaseModuleContext)
+
+	// ProcessBazelQueryResponse uses Bazel information (obtained from the BazelContext)
+	// to set module fields and providers to propagate this module's metadata upstream.
+	// This effectively "bridges the gap" between Bazel and Soong in a mixed build.
+	// Soong modules depending on this module should be oblivious to the fact that
+	// this module was handled by Bazel.
+	ProcessBazelQueryResponse(ctx ModuleContext)
+}
+
 // BazelModule is a lightweight wrapper interface around Module for Bazel-convertible modules.
 type BazelModule interface {
 	Module
@@ -342,7 +363,7 @@
 // converted or handcrafted Bazel target. As a side effect, calling this
 // method will also log whether this module is mixed build enabled for
 // metrics reporting.
-func MixedBuildsEnabled(ctx ModuleContext) bool {
+func MixedBuildsEnabled(ctx BaseModuleContext) bool {
 	mixedBuildEnabled := mixedBuildPossible(ctx)
 	ctx.Config().LogMixedBuild(ctx, mixedBuildEnabled)
 	return mixedBuildEnabled
@@ -350,7 +371,7 @@
 
 // mixedBuildPossible returns true if a module is ready to be replaced by a
 // converted or handcrafted Bazel target.
-func mixedBuildPossible(ctx ModuleContext) bool {
+func mixedBuildPossible(ctx BaseModuleContext) bool {
 	if ctx.Os() == Windows {
 		// Windows toolchains are not currently supported.
 		return false
diff --git a/android/bazel_handler.go b/android/bazel_handler.go
index c4eb0f3..4d9423a 100644
--- a/android/bazel_handler.go
+++ b/android/bazel_handler.go
@@ -33,6 +33,26 @@
 	"android/soong/bazel"
 )
 
+func init() {
+	RegisterMixedBuildsMutator(InitRegistrationContext)
+}
+
+func RegisterMixedBuildsMutator(ctx RegistrationContext) {
+	ctx.PostDepsMutators(func(ctx RegisterMutatorsContext) {
+		ctx.BottomUp("mixed_builds_prep", mixedBuildsPrepareMutator).Parallel()
+	})
+}
+
+func mixedBuildsPrepareMutator(ctx BottomUpMutatorContext) {
+	if m := ctx.Module(); m.Enabled() {
+		if mixedBuildMod, ok := m.(MixedBuildBuildable); ok {
+			if mixedBuildMod.IsMixedBuildSupported(ctx) && MixedBuildsEnabled(ctx) {
+				mixedBuildMod.QueueBazelCall(ctx)
+			}
+		}
+	}
+}
+
 type cqueryRequest interface {
 	// Name returns a string name for this request type. Such request type names must be unique,
 	// and must only consist of alphanumeric characters.
@@ -62,33 +82,32 @@
 	configKey   configKey
 }
 
-// bazelHandler is the interface for a helper object related to deferring to Bazel for
-// processing a module (during Bazel mixed builds). Individual module types should define
-// their own bazel handler if they support deferring to Bazel.
-type BazelHandler interface {
-	// Issue query to Bazel to retrieve information about Bazel's view of the current module.
-	// If Bazel returns this information, set module properties on the current module to reflect
-	// the returned information.
-	// Returns true if information was available from Bazel, false if bazel invocation still needs to occur.
-	GenerateBazelBuildActions(ctx ModuleContext, label string) bool
-}
-
+// BazelContext is a context object useful for interacting with Bazel during
+// the course of a build. Use of Bazel to evaluate part of the build graph
+// is referred to as a "mixed build". (Some modules are managed by Soong,
+// some are managed by Bazel). To facilitate interop between these build
+// subgraphs, Soong may make requests to Bazel and evaluate their responses
+// so that Soong modules may accurately depend on Bazel targets.
 type BazelContext interface {
-	// The methods below involve queuing cquery requests to be later invoked
-	// by bazel. If any of these methods return (_, false), then the request
-	// has been queued to be run later.
+	// Add a cquery request to the bazel request queue. All queued requests
+	// will be sent to Bazel on a subsequent invocation of InvokeBazel.
+	QueueBazelRequest(label string, requestType cqueryRequest, cfgKey configKey)
+
+	// ** Cquery Results Retrieval Functions
+	// The below functions pertain to retrieving cquery results from a prior
+	// InvokeBazel function call and parsing the results.
 
 	// Returns result files built by building the given bazel target label.
-	GetOutputFiles(label string, cfgKey configKey) ([]string, bool)
+	GetOutputFiles(label string, cfgKey configKey) ([]string, error)
 
-	// TODO(cparsons): Other cquery-related methods should be added here.
 	// Returns the results of GetOutputFiles and GetCcObjectFiles in a single query (in that order).
-	GetCcInfo(label string, cfgKey configKey) (cquery.CcInfo, bool, error)
+	GetCcInfo(label string, cfgKey configKey) (cquery.CcInfo, error)
 
 	// Returns the executable binary resultant from building together the python sources
-	GetPythonBinary(label string, cfgKey configKey) (string, bool)
+	// TODO(b/232976601): Remove.
+	GetPythonBinary(label string, cfgKey configKey) (string, error)
 
-	// ** End cquery methods
+	// ** end Cquery Results Retrieval Functions
 
 	// Issues commands to Bazel to receive results for all cquery requests
 	// queued in the BazelContext.
@@ -153,19 +172,23 @@
 	LabelToPythonBinary map[string]string
 }
 
-func (m MockBazelContext) GetOutputFiles(label string, cfgKey configKey) ([]string, bool) {
-	result, ok := m.LabelToOutputFiles[label]
-	return result, ok
+func (m MockBazelContext) QueueBazelRequest(label string, requestType cqueryRequest, cfgKey configKey) {
+	panic("unimplemented")
 }
 
-func (m MockBazelContext) GetCcInfo(label string, cfgKey configKey) (cquery.CcInfo, bool, error) {
-	result, ok := m.LabelToCcInfo[label]
-	return result, ok, nil
+func (m MockBazelContext) GetOutputFiles(label string, cfgKey configKey) ([]string, error) {
+	result, _ := m.LabelToOutputFiles[label]
+	return result, nil
 }
 
-func (m MockBazelContext) GetPythonBinary(label string, cfgKey configKey) (string, bool) {
-	result, ok := m.LabelToPythonBinary[label]
-	return result, ok
+func (m MockBazelContext) GetCcInfo(label string, cfgKey configKey) (cquery.CcInfo, error) {
+	result, _ := m.LabelToCcInfo[label]
+	return result, nil
+}
+
+func (m MockBazelContext) GetPythonBinary(label string, cfgKey configKey) (string, error) {
+	result, _ := m.LabelToPythonBinary[label]
+	return result, nil
 }
 
 func (m MockBazelContext) InvokeBazel() error {
@@ -188,46 +211,53 @@
 
 var _ BazelContext = MockBazelContext{}
 
-func (bazelCtx *bazelContext) GetOutputFiles(label string, cfgKey configKey) ([]string, bool) {
-	rawString, ok := bazelCtx.cquery(label, cquery.GetOutputFiles, cfgKey)
-	var ret []string
-	if ok {
+func (bazelCtx *bazelContext) QueueBazelRequest(label string, requestType cqueryRequest, cfgKey configKey) {
+	key := cqueryKey{label, requestType, cfgKey}
+	bazelCtx.requestMutex.Lock()
+	defer bazelCtx.requestMutex.Unlock()
+	bazelCtx.requests[key] = true
+}
+
+func (bazelCtx *bazelContext) GetOutputFiles(label string, cfgKey configKey) ([]string, error) {
+	key := cqueryKey{label, cquery.GetOutputFiles, cfgKey}
+	if rawString, ok := bazelCtx.results[key]; ok {
 		bazelOutput := strings.TrimSpace(rawString)
-		ret = cquery.GetOutputFiles.ParseResult(bazelOutput)
+		return cquery.GetOutputFiles.ParseResult(bazelOutput), nil
 	}
-	return ret, ok
+	return nil, fmt.Errorf("no bazel response found for %v", key)
 }
 
-func (bazelCtx *bazelContext) GetCcInfo(label string, cfgKey configKey) (cquery.CcInfo, bool, error) {
-	result, ok := bazelCtx.cquery(label, cquery.GetCcInfo, cfgKey)
-	if !ok {
-		return cquery.CcInfo{}, ok, nil
-	}
-
-	bazelOutput := strings.TrimSpace(result)
-	ret, err := cquery.GetCcInfo.ParseResult(bazelOutput)
-	return ret, ok, err
-}
-
-func (bazelCtx *bazelContext) GetPythonBinary(label string, cfgKey configKey) (string, bool) {
-	rawString, ok := bazelCtx.cquery(label, cquery.GetPythonBinary, cfgKey)
-	var ret string
-	if ok {
+func (bazelCtx *bazelContext) GetCcInfo(label string, cfgKey configKey) (cquery.CcInfo, error) {
+	key := cqueryKey{label, cquery.GetCcInfo, cfgKey}
+	if rawString, ok := bazelCtx.results[key]; ok {
 		bazelOutput := strings.TrimSpace(rawString)
-		ret = cquery.GetPythonBinary.ParseResult(bazelOutput)
+		return cquery.GetCcInfo.ParseResult(bazelOutput)
 	}
-	return ret, ok
+	return cquery.CcInfo{}, fmt.Errorf("no bazel response found for %v", key)
 }
 
-func (n noopBazelContext) GetOutputFiles(label string, cfgKey configKey) ([]string, bool) {
+func (bazelCtx *bazelContext) GetPythonBinary(label string, cfgKey configKey) (string, error) {
+	key := cqueryKey{label, cquery.GetPythonBinary, cfgKey}
+	if rawString, ok := bazelCtx.results[key]; ok {
+		bazelOutput := strings.TrimSpace(rawString)
+		return cquery.GetPythonBinary.ParseResult(bazelOutput), nil
+	}
+	return "", fmt.Errorf("no bazel response found for %v", key)
+}
+
+func (n noopBazelContext) QueueBazelRequest(label string, requestType cqueryRequest, cfgKey configKey) {
 	panic("unimplemented")
 }
 
-func (n noopBazelContext) GetCcInfo(label string, cfgKey configKey) (cquery.CcInfo, bool, error) {
+func (n noopBazelContext) GetOutputFiles(label string, cfgKey configKey) ([]string, error) {
 	panic("unimplemented")
 }
 
-func (n noopBazelContext) GetPythonBinary(label string, cfgKey configKey) (string, bool) {
+func (n noopBazelContext) GetCcInfo(label string, cfgKey configKey) (cquery.CcInfo, error) {
+	panic("unimplemented")
+}
+
+func (n noopBazelContext) GetPythonBinary(label string, cfgKey configKey) (string, error) {
 	panic("unimplemented")
 }
 
@@ -314,24 +344,6 @@
 	return true
 }
 
-// Adds a cquery request to the Bazel request queue, to be later invoked, or
-// returns the result of the given request if the request was already made.
-// 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 cqueryRequest,
-	cfgKey configKey) (string, bool) {
-	key := cqueryKey{label, requestType, cfgKey}
-	if result, ok := context.results[key]; ok {
-		return result, true
-	} else {
-		context.requestMutex.Lock()
-		defer context.requestMutex.Unlock()
-		context.requests[key] = true
-		return "", false
-	}
-}
-
 func pwdPrefix() string {
 	// Darwin doesn't have /proc
 	if runtime.GOOS != "darwin" {
@@ -916,7 +928,7 @@
 	return arch + "|" + os
 }
 
-func GetConfigKey(ctx ModuleContext) configKey {
+func GetConfigKey(ctx BaseModuleContext) configKey {
 	return configKey{
 		// use string because Arch is not a valid key in go
 		arch:   ctx.Arch().String(),
diff --git a/android/bazel_handler_test.go b/android/bazel_handler_test.go
index e5cff90..cfdccd7 100644
--- a/android/bazel_handler_test.go
+++ b/android/bazel_handler_test.go
@@ -5,6 +5,8 @@
 	"path/filepath"
 	"reflect"
 	"testing"
+
+	"android/soong/bazel/cquery"
 )
 
 func TestRequestResultsAfterInvokeBazel(t *testing.T) {
@@ -13,17 +15,14 @@
 	bazelContext, _ := testBazelContext(t, map[bazelCommand]string{
 		bazelCommand{command: "cquery", expression: "deps(@soong_injection//mixed_builds:buildroot, 2)"}: `//foo:bar|arm64_armv8-a|android>>out/foo/bar.txt`,
 	})
-	g, ok := bazelContext.GetOutputFiles(label, cfg)
-	if ok {
-		t.Errorf("Did not expect cquery results prior to running InvokeBazel(), but got %s", g)
-	}
+	bazelContext.QueueBazelRequest(label, cquery.GetOutputFiles, cfg)
 	err := bazelContext.InvokeBazel()
 	if err != nil {
 		t.Fatalf("Did not expect error invoking Bazel, but got %s", err)
 	}
-	g, ok = bazelContext.GetOutputFiles(label, cfg)
-	if !ok {
-		t.Errorf("Expected cquery results after running InvokeBazel(), but got none")
+	g, err := bazelContext.GetOutputFiles(label, cfg)
+	if err != nil {
+		t.Errorf("Expected cquery results after running InvokeBazel(), but got err %v", err)
 	} else if w := []string{"out/foo/bar.txt"}; !reflect.DeepEqual(w, g) {
 		t.Errorf("Expected output %s, got %s", w, g)
 	}
diff --git a/android/config.go b/android/config.go
index d695217..9e96c6a 100644
--- a/android/config.go
+++ b/android/config.go
@@ -2047,7 +2047,7 @@
 	return Bool(c.productVariables.HostMusl)
 }
 
-func (c *config) LogMixedBuild(ctx ModuleContext, useBazel bool) {
+func (c *config) LogMixedBuild(ctx BaseModuleContext, useBazel bool) {
 	moduleName := ctx.Module().Name()
 	c.mixedBuildsLock.Lock()
 	defer c.mixedBuildsLock.Unlock()
diff --git a/android/filegroup.go b/android/filegroup.go
index 1bf5e07..485e0b9 100644
--- a/android/filegroup.go
+++ b/android/filegroup.go
@@ -18,6 +18,7 @@
 	"strings"
 
 	"android/soong/bazel"
+	"android/soong/bazel/cquery"
 
 	"github.com/google/blueprint"
 )
@@ -101,6 +102,7 @@
 	srcs       Paths
 }
 
+var _ MixedBuildBuildable = (*fileGroup)(nil)
 var _ SourceFileProducer = (*fileGroup)(nil)
 
 // filegroup contains a list of files that are referenced by other modules
@@ -114,42 +116,11 @@
 	return module
 }
 
-func (fg *fileGroup) maybeGenerateBazelBuildActions(ctx ModuleContext) {
-	if !MixedBuildsEnabled(ctx) {
-		return
-	}
-
-	archVariant := ctx.Arch().String()
-	osVariant := ctx.Os()
-	if len(fg.Srcs()) == 1 && fg.Srcs()[0].Base() == fg.Name() {
-		// This will be a regular file target, not filegroup, in Bazel.
-		// See FilegroupBp2Build for more information.
-		archVariant = Common.String()
-		osVariant = CommonOS
-	}
-
-	bazelCtx := ctx.Config().BazelContext
-	filePaths, ok := bazelCtx.GetOutputFiles(fg.GetBazelLabel(ctx, fg), configKey{archVariant, osVariant})
-	if !ok {
-		return
-	}
-
-	bazelOuts := make(Paths, 0, len(filePaths))
-	for _, p := range filePaths {
-		src := PathForBazelOut(ctx, p)
-		bazelOuts = append(bazelOuts, src)
-	}
-
-	fg.srcs = bazelOuts
-}
-
 func (fg *fileGroup) GenerateAndroidBuildActions(ctx ModuleContext) {
 	fg.srcs = PathsForModuleSrcExcludes(ctx, fg.properties.Srcs, fg.properties.Exclude_srcs)
 	if fg.properties.Path != nil {
 		fg.srcs = PathsWithModuleSrcSubDir(ctx, fg.srcs, String(fg.properties.Path))
 	}
-
-	fg.maybeGenerateBazelBuildActions(ctx)
 }
 
 func (fg *fileGroup) Srcs() Paths {
@@ -161,3 +132,38 @@
 		ctx.StrictRaw(makeVar, strings.Join(fg.srcs.Strings(), " "))
 	}
 }
+
+func (fg *fileGroup) QueueBazelCall(ctx BaseModuleContext) {
+	bazelCtx := ctx.Config().BazelContext
+
+	bazelCtx.QueueBazelRequest(
+		fg.GetBazelLabel(ctx, fg),
+		cquery.GetOutputFiles,
+		configKey{Common.String(), CommonOS})
+}
+
+func (fg *fileGroup) IsMixedBuildSupported(ctx BaseModuleContext) bool {
+	return true
+}
+
+func (fg *fileGroup) ProcessBazelQueryResponse(ctx ModuleContext) {
+	fg.srcs = PathsForModuleSrcExcludes(ctx, fg.properties.Srcs, fg.properties.Exclude_srcs)
+	if fg.properties.Path != nil {
+		fg.srcs = PathsWithModuleSrcSubDir(ctx, fg.srcs, String(fg.properties.Path))
+	}
+
+	bazelCtx := ctx.Config().BazelContext
+	filePaths, err := bazelCtx.GetOutputFiles(fg.GetBazelLabel(ctx, fg), configKey{Common.String(), CommonOS})
+	if err != nil {
+		ctx.ModuleErrorf(err.Error())
+		return
+	}
+
+	bazelOuts := make(Paths, 0, len(filePaths))
+	for _, p := range filePaths {
+		src := PathForBazelOut(ctx, p)
+		bazelOuts = append(bazelOuts, src)
+	}
+
+	fg.srcs = bazelOuts
+}
diff --git a/android/module.go b/android/module.go
index ab68e24..97a46dd 100644
--- a/android/module.go
+++ b/android/module.go
@@ -2275,7 +2275,11 @@
 			return
 		}
 
-		m.module.GenerateAndroidBuildActions(ctx)
+		if mixedBuildMod, handled := m.isHandledByBazel(ctx); handled {
+			mixedBuildMod.ProcessBazelQueryResponse(ctx)
+		} else {
+			m.module.GenerateAndroidBuildActions(ctx)
+		}
 		if ctx.Failed() {
 			return
 		}
@@ -2331,6 +2335,18 @@
 	m.variables = ctx.variables
 }
 
+func (m *ModuleBase) isHandledByBazel(ctx ModuleContext) (MixedBuildBuildable, bool) {
+	if !ctx.Config().BazelContext.BazelEnabled() {
+		return nil, false
+	}
+	if mixedBuildMod, ok := m.module.(MixedBuildBuildable); ok {
+		if mixedBuildMod.IsMixedBuildSupported(ctx) && MixedBuildsEnabled(ctx) {
+			return mixedBuildMod, true
+		}
+	}
+	return nil, false
+}
+
 // Check the supplied dist structure to make sure that it is valid.
 //
 // property - the base property, e.g. dist or dists[1], which is combined with the