Reimplement afdo support for rust

Ignore-AOSP-First: The parent CL is internal
Bug: 267229065
Test: go test
Change-Id: Ia14679285b92f3f14ff269392a61f978c71311b2
Merged-In: Ia14679285b92f3f14ff269392a61f978c71311b2

rust/afdo.go

@@ -17,7 +17,10 @@ package rust
 import (
 	"fmt"
 
+	"android/soong/android"
 	"android/soong/cc"
+
+	"github.com/google/blueprint"
 )
 
 const afdoFlagFormat = "-Zprofile-sample-use=%s"
@@ -30,19 +33,49 @@ func (afdo *afdo) props() []interface{} {
 	return []interface{}{&afdo.Properties}
 }
 
-func (afdo *afdo) flags(ctx ModuleContext, flags Flags, deps PathDeps) (Flags, PathDeps) {
+func (afdo *afdo) addDep(ctx BaseModuleContext, actx android.BottomUpMutatorContext) {
+	// afdo is not supported outside of Android
+	if ctx.Host() {
+		return
+	}
+
+	if mod, ok := ctx.Module().(*Module); ok && mod.Enabled() {
+		fdoProfileName, err := actx.DeviceConfig().AfdoProfile(actx.ModuleName())
+		if err != nil {
+			ctx.ModuleErrorf("%s", err.Error())
+		}
+		if fdoProfileName != nil {
+			actx.AddFarVariationDependencies(
+				[]blueprint.Variation{
+					{Mutator: "arch", Variation: actx.Target().ArchVariation()},
+					{Mutator: "os", Variation: "android"},
+				},
+				cc.FdoProfileTag,
+				[]string{*fdoProfileName}...,
+			)
+		}
+	}
+}
+
+func (afdo *afdo) flags(ctx android.ModuleContext, flags Flags, deps PathDeps) (Flags, PathDeps) {
 	if ctx.Host() {
 		return flags, deps
 	}
 
-	if afdo != nil && afdo.Properties.Afdo {
-		if profileFile := afdo.Properties.GetAfdoProfileFile(ctx, ctx.ModuleName()); profileFile.Valid() {
-			profileUseFlag := fmt.Sprintf(afdoFlagFormat, profileFile)
+	if !afdo.Properties.Afdo {
+		return flags, deps
+	}
+
+	ctx.VisitDirectDepsWithTag(cc.FdoProfileTag, func(m android.Module) {
+		if ctx.OtherModuleHasProvider(m, cc.FdoProfileProvider) {
+			info := ctx.OtherModuleProvider(m, cc.FdoProfileProvider).(cc.FdoProfileInfo)
+			path := info.Path
+			profileUseFlag := fmt.Sprintf(afdoFlagFormat, path.String())
 			flags.RustFlags = append(flags.RustFlags, profileUseFlag)
 
-			profileFilePath := profileFile.Path()
-			deps.AfdoProfiles = append(deps.AfdoProfiles, profileFilePath)
+			deps.AfdoProfiles = append(deps.AfdoProfiles, path)
 		}
-	}
+	})
+
 	return flags, deps
 }

rust/afdo_test.go

@@ -16,6 +16,7 @@ package rust
 
 import (
 	"android/soong/android"
+	"android/soong/cc"
 	"fmt"
 	"strings"
 	"testing"
@@ -31,13 +32,27 @@ func TestAfdoEnabled(t *testing.T) {
 `
 	result := android.GroupFixturePreparers(
 		prepareForRustTest,
-		android.FixtureAddTextFile("toolchain/pgo-profiles/sampling/foo.afdo", ""),
+		cc.PrepareForTestWithFdoProfile,
+		android.FixtureAddTextFile("afdo_profiles_package/foo.afdo", ""),
+		android.FixtureModifyProductVariables(func(variables android.FixtureProductVariables) {
+			variables.AfdoProfiles = []string{
+				"foo://afdo_profiles_package:foo_afdo",
+			}
+		}),
+		android.MockFS{
+			"afdo_profiles_package/Android.bp": []byte(`
+				fdo_profile {
+					name: "foo_afdo",
+					profile: "foo.afdo",
+				}
+			`),
+		}.AddToFixture(),
 		rustMockedFiles.AddToFixture(),
 	).RunTestWithBp(t, bp)
 
 	foo := result.ModuleForTests("foo", "android_arm64_armv8-a").Rule("rustc")
 
-	expectedCFlag := fmt.Sprintf(afdoFlagFormat, "toolchain/pgo-profiles/sampling/foo.afdo")
+	expectedCFlag := fmt.Sprintf(afdoFlagFormat, "afdo_profiles_package/foo.afdo")
 
 	if !strings.Contains(foo.Args["rustcFlags"], expectedCFlag) {
 		t.Errorf("Expected 'foo' to enable afdo, but did not find %q in cflags %q", expectedCFlag, foo.Args["rustcFlags"])
@@ -55,16 +70,37 @@ func TestAfdoEnabledWithMultiArchs(t *testing.T) {
 `
 	result := android.GroupFixturePreparers(
 		prepareForRustTest,
-		android.FixtureAddTextFile("toolchain/pgo-profiles/sampling/foo_arm.afdo", ""),
-		android.FixtureAddTextFile("toolchain/pgo-profiles/sampling/foo_arm64.afdo", ""),
+		cc.PrepareForTestWithFdoProfile,
+		android.FixtureAddTextFile("afdo_profiles_package/foo_arm.afdo", ""),
+		android.FixtureAddTextFile("afdo_profiles_package/foo_arm64.afdo", ""),
+		android.FixtureModifyProductVariables(func(variables android.FixtureProductVariables) {
+			variables.AfdoProfiles = []string{
+				"foo://afdo_profiles_package:foo_afdo",
+			}
+		}),
+		android.MockFS{
+			"afdo_profiles_package/Android.bp": []byte(`
+				fdo_profile {
+					name: "foo_afdo",
+					arch: {
+						arm: {
+							profile: "foo_arm.afdo",
+						},
+						arm64: {
+							profile: "foo_arm64.afdo",
+						}
+					}
+				}
+			`),
+		}.AddToFixture(),
 		rustMockedFiles.AddToFixture(),
 	).RunTestWithBp(t, bp)
 
 	fooArm := result.ModuleForTests("foo", "android_arm_armv7-a-neon").Rule("rustc")
 	fooArm64 := result.ModuleForTests("foo", "android_arm64_armv8-a").Rule("rustc")
 
-	expectedCFlagArm := fmt.Sprintf(afdoFlagFormat, "toolchain/pgo-profiles/sampling/foo_arm.afdo")
-	expectedCFlagArm64 := fmt.Sprintf(afdoFlagFormat, "toolchain/pgo-profiles/sampling/foo_arm64.afdo")
+	expectedCFlagArm := fmt.Sprintf(afdoFlagFormat, "afdo_profiles_package/foo_arm.afdo")
+	expectedCFlagArm64 := fmt.Sprintf(afdoFlagFormat, "afdo_profiles_package/foo_arm64.afdo")
 
 	if !strings.Contains(fooArm.Args["rustcFlags"], expectedCFlagArm) {
 		t.Errorf("Expected 'fooArm' to enable afdo, but did not find %q in cflags %q", expectedCFlagArm, fooArm.Args["rustcFlags"])

rust/rust.go

@@ -39,7 +39,6 @@ func init() {
 		ctx.BottomUp("rust_libraries", LibraryMutator).Parallel()
 		ctx.BottomUp("rust_stdlinkage", LibstdMutator).Parallel()
 		ctx.BottomUp("rust_begin", BeginMutator).Parallel()
-
 	})
 	android.PostDepsMutators(func(ctx android.RegisterMutatorsContext) {
 		ctx.BottomUp("rust_sanitizers", rustSanitizerRuntimeMutator).Parallel()
@@ -920,7 +919,7 @@ func (mod *Module) GenerateAndroidBuildActions(actx android.ModuleContext) {
 
 	// Calculate rustc flags
 	if mod.afdo != nil {
-		flags, deps = mod.afdo.flags(ctx, flags, deps)
+		flags, deps = mod.afdo.flags(actx, flags, deps)
 	}
 	if mod.compiler != nil {
 		flags = mod.compiler.compilerFlags(ctx, flags)
@@ -1613,6 +1612,8 @@ func (mod *Module) DepsMutator(actx android.BottomUpMutatorContext) {
 
 	// proc_macros are compiler plugins, and so we need the host arch variant as a dependendcy.
 	actx.AddFarVariationDependencies(ctx.Config().BuildOSTarget.Variations(), procMacroDepTag, deps.ProcMacros...)
+
+	mod.afdo.addDep(ctx, actx)
 }
 
 // addRlibDependency will add an rlib dependency, rewriting to the snapshot library if available.