Fix genrule tool dependencies when a prebuilt tool is preferred.

Since the genrule tool dep mutator runs after the prebuilt mutators,
this adds a helper function android.PrebuiltGetPreferred to resolve the
source or prebuilt as appropriate.

Test: m SOONG_CONFIG_art_module_source_build=false droid
  in internal
Bug: 214292395
Change-Id: I1a208fd048b998f9f19ad1f45d8389decda2cb9e

android/prebuilt.go

@@ -309,6 +309,54 @@ func GetEmbeddedPrebuilt(module Module) *Prebuilt {
 	return nil
 }
 
+// PrebuiltGetPreferred returns the module that is preferred for the given
+// module. That is either the module itself or the prebuilt counterpart that has
+// taken its place. The given module must be a direct dependency of the current
+// context module, and it must be the source module if both source and prebuilt
+// exist.
+//
+// This function is for use on dependencies after PrebuiltPostDepsMutator has
+// run - any dependency that is registered before that will already reference
+// the right module. This function is only safe to call after all mutators that
+// may call CreateVariations, e.g. in GenerateAndroidBuildActions.
+func PrebuiltGetPreferred(ctx BaseModuleContext, module Module) Module {
+	if !module.IsReplacedByPrebuilt() {
+		return module
+	}
+	if IsModulePrebuilt(module) {
+		// If we're given a prebuilt then assume there's no source module around.
+		return module
+	}
+
+	sourceModDepFound := false
+	var prebuiltMod Module
+
+	ctx.WalkDeps(func(child, parent Module) bool {
+		if prebuiltMod != nil {
+			return false
+		}
+		if parent == ctx.Module() {
+			// First level: Only recurse if the module is found as a direct dependency.
+			sourceModDepFound = child == module
+			return sourceModDepFound
+		}
+		// Second level: Follow PrebuiltDepTag to the prebuilt.
+		if t := ctx.OtherModuleDependencyTag(child); t == PrebuiltDepTag {
+			prebuiltMod = child
+		}
+		return false
+	})
+
+	if prebuiltMod == nil {
+		if !sourceModDepFound {
+			panic(fmt.Errorf("Failed to find source module as a direct dependency: %s", module))
+		} else {
+			panic(fmt.Errorf("Failed to find prebuilt for source module: %s", module))
+		}
+	}
+	return prebuiltMod
+}
+
 func RegisterPrebuiltsPreArchMutators(ctx RegisterMutatorsContext) {
 	ctx.BottomUp("prebuilt_rename", PrebuiltRenameMutator).Parallel()
 }

genrule/genrule.go

@@ -106,6 +106,16 @@ type hostToolDependencyTag struct {
 	android.LicenseAnnotationToolchainDependencyTag
 	label string
 }
+
+func (t hostToolDependencyTag) AllowDisabledModuleDependency(target android.Module) bool {
+	// Allow depending on a disabled module if it's replaced by a prebuilt
+	// counterpart. We get the prebuilt through android.PrebuiltGetPreferred in
+	// GenerateAndroidBuildActions.
+	return target.IsReplacedByPrebuilt()
+}
+
+var _ android.AllowDisabledModuleDependency = (*hostToolDependencyTag)(nil)
+
 type generatorProperties struct {
 	// The command to run on one or more input files. Cmd supports substitution of a few variables.
 	//
@@ -298,6 +308,12 @@ func (g *Module) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 			switch tag := ctx.OtherModuleDependencyTag(module).(type) {
 			case hostToolDependencyTag:
 				tool := ctx.OtherModuleName(module)
+				if m, ok := module.(android.Module); ok {
+					// Necessary to retrieve any prebuilt replacement for the tool, since
+					// toolDepsMutator runs too late for the prebuilt mutators to have
+					// replaced the dependency.
+					module = android.PrebuiltGetPreferred(ctx, m)
+				}
 
 				switch t := module.(type) {
 				case android.HostToolProvider:

genrule/genrule_test.go

@@ -34,7 +34,9 @@ var prepareForGenRuleTest = android.GroupFixturePreparers(
 	android.PrepareForTestWithFilegroup,
 	PrepareForTestWithGenRuleBuildComponents,
 	android.FixtureRegisterWithContext(func(ctx android.RegistrationContext) {
+		android.RegisterPrebuiltMutators(ctx)
 		ctx.RegisterModuleType("tool", toolFactory)
+		ctx.RegisterModuleType("prebuilt_tool", prebuiltToolFactory)
 		ctx.RegisterModuleType("output", outputProducerFactory)
 		ctx.RegisterModuleType("use_source", useSourceFactory)
 	}),
@@ -720,6 +722,69 @@ func TestGenruleOutputFiles(t *testing.T) {
 		result.ModuleForTests("gen_all", "").Module().(*useSource).srcs)
 }
 
+func TestPrebuiltTool(t *testing.T) {
+	testcases := []struct {
+		name             string
+		bp               string
+		expectedToolName string
+	}{
+		{
+			name: "source only",
+			bp: `
+				tool { name: "tool" }
+			`,
+			expectedToolName: "bin/tool",
+		},
+		{
+			name: "prebuilt only",
+			bp: `
+				prebuilt_tool { name: "tool" }
+			`,
+			expectedToolName: "prebuilt_bin/tool",
+		},
+		{
+			name: "source preferred",
+			bp: `
+				tool { name: "tool" }
+				prebuilt_tool { name: "tool" }
+			`,
+			expectedToolName: "bin/tool",
+		},
+		{
+			name: "prebuilt preferred",
+			bp: `
+				tool { name: "tool" }
+				prebuilt_tool { name: "tool", prefer: true }
+			`,
+			expectedToolName: "prebuilt_bin/prebuilt_tool",
+		},
+		{
+			name: "source disabled",
+			bp: `
+				tool { name: "tool", enabled: false }
+				prebuilt_tool { name: "tool" }
+      `,
+			expectedToolName: "prebuilt_bin/prebuilt_tool",
+		},
+	}
+
+	for _, test := range testcases {
+		t.Run(test.name, func(t *testing.T) {
+			result := prepareForGenRuleTest.RunTestWithBp(t, test.bp+`
+				genrule {
+					name: "gen",
+					tools: ["tool"],
+					out: ["foo"],
+					cmd: "$(location tool)",
+				}
+			`)
+			gen := result.Module("gen", "").(*Module)
+			expectedCmd := "__SBOX_SANDBOX_DIR__/tools/out/" + test.expectedToolName
+			android.AssertStringEquals(t, "command", expectedCmd, gen.rawCommands[0])
+		})
+	}
+}
+
 func TestGenruleWithBazel(t *testing.T) {
 	bp := `
 		genrule {
@@ -764,7 +829,33 @@ func (t *testTool) HostToolPath() android.OptionalPath {
 	return android.OptionalPathForPath(t.outputFile)
 }
 
+type prebuiltTestTool struct {
+	android.ModuleBase
+	prebuilt android.Prebuilt
+	testTool
+}
+
+func (p *prebuiltTestTool) Name() string {
+	return p.prebuilt.Name(p.ModuleBase.Name())
+}
+
+func (p *prebuiltTestTool) Prebuilt() *android.Prebuilt {
+	return &p.prebuilt
+}
+
+func (t *prebuiltTestTool) GenerateAndroidBuildActions(ctx android.ModuleContext) {
+	t.outputFile = ctx.InstallFile(android.PathForModuleInstall(ctx, "prebuilt_bin"), ctx.ModuleName(), android.PathForOutput(ctx, ctx.ModuleName()))
+}
+
+func prebuiltToolFactory() android.Module {
+	module := &prebuiltTestTool{}
+	android.InitPrebuiltModuleWithoutSrcs(module)
+	android.InitAndroidArchModule(module, android.HostSupported, android.MultilibFirst)
+	return module
+}
+
 var _ android.HostToolProvider = (*testTool)(nil)
+var _ android.HostToolProvider = (*prebuiltTestTool)(nil)
 
 type testOutputProducer struct {
 	android.ModuleBase