Use r8/d8 optimized profile for dexpreopt

Currently, dexpreopt supports profile guided optimization. This does not
work well with r8/d8 optimization, since the checked-in profile will not
match the dex signatures after r8/d8 has optimized the dex code.

This CL introduces a new property `dex_preopt.enable_profile_rewrting`.
If set, the checked-in profile will passed as `input` to r8 via
`--art-profile <input> <output>`. The <output> from the previous command
will be used as the profile for dexpreopt.

Test: m nothing --no-skip-soong-tests
Test: m CredentialManager with https://ag.corp.google.com/27448930
and obfuscation turned on
Test: nm -U symbol.odex # contains obfuscated methods

Bug: 335418838

Change-Id: I53beed9ed76f013262f1c503de0f2b74997c2a7f

java/base.go

@@ -1650,11 +1650,23 @@ func (j *Module) compile(ctx android.ModuleContext, extraSrcJars, extraClasspath
 				classesJar:    implementationAndResourcesJar,
 				jarName:       jarName,
 			}
-			dexOutputFile = j.dexer.compileDex(ctx, params)
+			if j.EnableProfileRewriting() {
+				profile := j.GetProfile()
+				if profile == "" || !j.GetProfileGuided() {
+					ctx.PropertyErrorf("enable_profile_rewriting", "Profile and Profile_guided must be set when enable_profile_rewriting is true")
+				}
+				params.artProfileInput = &profile
+			}
+			dexOutputFile, dexArtProfileOutput := j.dexer.compileDex(ctx, params)
 			if ctx.Failed() {
 				return
 			}
 
+			// If r8/d8 provides a profile that matches the optimized dex, use that for dexpreopt.
+			if dexArtProfileOutput != nil {
+				j.dexpreopter.SetRewrittenProfile(*dexArtProfileOutput)
+			}
+
 			// merge dex jar with resources if necessary
 			if j.resourceJar != nil {
 				jars := android.Paths{dexOutputFile, j.resourceJar}

java/dex.go

@@ -253,17 +253,25 @@ func (d *dexer) dexCommonFlags(ctx android.ModuleContext,
 	return flags, deps
 }
 
-func d8Flags(flags javaBuilderFlags) (d8Flags []string, d8Deps android.Paths) {
+func (d *dexer) d8Flags(ctx android.ModuleContext, dexParams *compileDexParams) (d8Flags []string, d8Deps android.Paths, artProfileOutput *android.OutputPath) {
+	flags := dexParams.flags
 	d8Flags = append(d8Flags, flags.bootClasspath.FormRepeatedClassPath("--lib ")...)
 	d8Flags = append(d8Flags, flags.dexClasspath.FormRepeatedClassPath("--lib ")...)
 
 	d8Deps = append(d8Deps, flags.bootClasspath...)
 	d8Deps = append(d8Deps, flags.dexClasspath...)
 
-	return d8Flags, d8Deps
+	if flags, deps, profileOutput := d.addArtProfile(ctx, dexParams); profileOutput != nil {
+		d8Flags = append(d8Flags, flags...)
+		d8Deps = append(d8Deps, deps...)
+		artProfileOutput = profileOutput
+	}
+
+	return d8Flags, d8Deps, artProfileOutput
 }
 
-func (d *dexer) r8Flags(ctx android.ModuleContext, flags javaBuilderFlags) (r8Flags []string, r8Deps android.Paths) {
+func (d *dexer) r8Flags(ctx android.ModuleContext, dexParams *compileDexParams) (r8Flags []string, r8Deps android.Paths, artProfileOutput *android.OutputPath) {
+	flags := dexParams.flags
 	opt := d.dexProperties.Optimize
 
 	// When an app contains references to APIs that are not in the SDK specified by
@@ -375,18 +383,44 @@ func (d *dexer) r8Flags(ctx android.ModuleContext, flags javaBuilderFlags) (r8Fl
 		}
 	}
 
-	return r8Flags, r8Deps
+	if flags, deps, profileOutput := d.addArtProfile(ctx, dexParams); profileOutput != nil {
+		r8Flags = append(r8Flags, flags...)
+		r8Deps = append(r8Deps, deps...)
+		artProfileOutput = profileOutput
+	}
+
+	return r8Flags, r8Deps, artProfileOutput
 }
 
 type compileDexParams struct {
-	flags         javaBuilderFlags
-	sdkVersion    android.SdkSpec
-	minSdkVersion android.ApiLevel
-	classesJar    android.Path
-	jarName       string
+	flags           javaBuilderFlags
+	sdkVersion      android.SdkSpec
+	minSdkVersion   android.ApiLevel
+	classesJar      android.Path
+	jarName         string
+	artProfileInput *string
 }
 
-func (d *dexer) compileDex(ctx android.ModuleContext, dexParams *compileDexParams) android.OutputPath {
+// Adds --art-profile to r8/d8 command.
+// r8/d8 will output a generated profile file to match the optimized dex code.
+func (d *dexer) addArtProfile(ctx android.ModuleContext, dexParams *compileDexParams) (flags []string, deps android.Paths, artProfileOutputPath *android.OutputPath) {
+	if dexParams.artProfileInput != nil {
+		artProfileInputPath := android.PathForModuleSrc(ctx, *dexParams.artProfileInput)
+		artProfileOutputPathValue := android.PathForModuleOut(ctx, "profile.prof.txt").OutputPath
+		artProfileOutputPath = &artProfileOutputPathValue
+		flags = []string{
+			"--art-profile",
+			artProfileInputPath.String(),
+			artProfileOutputPath.String(),
+		}
+		deps = append(deps, artProfileInputPath)
+	}
+	return flags, deps, artProfileOutputPath
+
+}
+
+// Return the compiled dex jar and (optional) profile _after_ r8 optimization
+func (d *dexer) compileDex(ctx android.ModuleContext, dexParams *compileDexParams) (android.OutputPath, *android.OutputPath) {
 
 	// Compile classes.jar into classes.dex and then javalib.jar
 	javalibJar := android.PathForModuleOut(ctx, "dex", dexParams.jarName).OutputPath
@@ -406,6 +440,7 @@ func (d *dexer) compileDex(ctx android.ModuleContext, dexParams *compileDexParam
 	}
 
 	useR8 := d.effectiveOptimizeEnabled()
+	var artProfileOutputPath *android.OutputPath
 	if useR8 {
 		proguardDictionary := android.PathForModuleOut(ctx, "proguard_dictionary")
 		d.proguardDictionary = android.OptionalPathForPath(proguardDictionary)
@@ -418,8 +453,19 @@ func (d *dexer) compileDex(ctx android.ModuleContext, dexParams *compileDexParam
 		d.proguardUsageZip = android.OptionalPathForPath(proguardUsageZip)
 		resourcesOutput := android.PathForModuleOut(ctx, "package-res-shrunken.apk")
 		d.resourcesOutput = android.OptionalPathForPath(resourcesOutput)
-		r8Flags, r8Deps := d.r8Flags(ctx, dexParams.flags)
-		r8Deps = append(r8Deps, commonDeps...)
+		implicitOutputs := android.WritablePaths{
+			proguardDictionary,
+			proguardUsageZip,
+			proguardConfiguration,
+		}
+		r8Flags, r8Deps, r8ArtProfileOutputPath := d.r8Flags(ctx, dexParams)
+		if r8ArtProfileOutputPath != nil {
+			artProfileOutputPath = r8ArtProfileOutputPath
+			implicitOutputs = append(
+				implicitOutputs,
+				artProfileOutputPath,
+			)
+		}
 		rule := r8
 		args := map[string]string{
 			"r8Flags":        strings.Join(append(commonFlags, r8Flags...), " "),
@@ -436,10 +482,6 @@ func (d *dexer) compileDex(ctx android.ModuleContext, dexParams *compileDexParam
 			rule = r8RE
 			args["implicits"] = strings.Join(r8Deps.Strings(), ",")
 		}
-		implicitOutputs := android.WritablePaths{
-			proguardDictionary,
-			proguardUsageZip,
-			proguardConfiguration}
 		if d.resourcesInput.Valid() {
 			implicitOutputs = append(implicitOutputs, resourcesOutput)
 			args["resourcesOutput"] = resourcesOutput.String()
@@ -454,18 +496,27 @@ func (d *dexer) compileDex(ctx android.ModuleContext, dexParams *compileDexParam
 			Args:            args,
 		})
 	} else {
-		d8Flags, d8Deps := d8Flags(dexParams.flags)
+		implicitOutputs := android.WritablePaths{}
+		d8Flags, d8Deps, d8ArtProfileOutputPath := d.d8Flags(ctx, dexParams)
+		if d8ArtProfileOutputPath != nil {
+			artProfileOutputPath = d8ArtProfileOutputPath
+			implicitOutputs = append(
+				implicitOutputs,
+				artProfileOutputPath,
+			)
+		}
 		d8Deps = append(d8Deps, commonDeps...)
 		rule := d8
 		if ctx.Config().UseRBE() && ctx.Config().IsEnvTrue("RBE_D8") {
 			rule = d8RE
 		}
 		ctx.Build(pctx, android.BuildParams{
-			Rule:        rule,
-			Description: "d8",
-			Output:      javalibJar,
-			Input:       dexParams.classesJar,
-			Implicits:   d8Deps,
+			Rule:            rule,
+			Description:     "d8",
+			Output:          javalibJar,
+			Input:           dexParams.classesJar,
+			ImplicitOutputs: implicitOutputs,
+			Implicits:       d8Deps,
 			Args: map[string]string{
 				"d8Flags":        strings.Join(append(commonFlags, d8Flags...), " "),
 				"zipFlags":       zipFlags,
@@ -480,5 +531,5 @@ func (d *dexer) compileDex(ctx android.ModuleContext, dexParams *compileDexParam
 		javalibJar = alignedJavalibJar
 	}
 
-	return javalibJar
+	return javalibJar, artProfileOutputPath
 }

java/dex_test.go

@@ -662,3 +662,30 @@ func TestProguardFlagsInheritanceAppImport(t *testing.T) {
 	android.AssertStringDoesContain(t, "expected aarimports's proguard flags",
 		appR8.Args["r8Flags"], "proguard.txt")
 }
+
+func TestR8FlagsArtProfile(t *testing.T) {
+	result := PrepareForTestWithJavaDefaultModules.RunTestWithBp(t, `
+		android_app {
+			name: "app",
+			srcs: ["foo.java"],
+			platform_apis: true,
+			dex_preopt: {
+				profile_guided: true,
+				profile: "profile.txt.prof",
+				enable_profile_rewriting: true,
+			},
+		}
+	`)
+
+	app := result.ModuleForTests("app", "android_common")
+	appR8 := app.Rule("r8")
+	android.AssertStringDoesContain(t, "expected --art-profile in app r8 flags",
+		appR8.Args["r8Flags"], "--art-profile")
+
+	appDexpreopt := app.Rule("dexpreopt")
+	android.AssertStringDoesContain(t,
+		"expected --art-profile output to be used to create .prof binary",
+		appDexpreopt.RuleParams.Command,
+		"--create-profile-from=out/soong/.intermediates/app/android_common/profile.prof.txt --output-profile-type=app",
+	)
+}

java/dexpreopt.go

@@ -19,6 +19,8 @@ import (
 	"sort"
 	"strings"
 
+	"github.com/google/blueprint/proptools"
+
 	"android/soong/android"
 	"android/soong/dexpreopt"
 )
@@ -139,6 +141,10 @@ type dexpreopter struct {
 	// The path to the profile that dexpreopter accepts. It must be in the binary format. If this is
 	// set, it overrides the profile settings in `dexpreoptProperties`.
 	inputProfilePathOnHost android.Path
+
+	// The path to the profile that matches the dex optimized by r8/d8. It is in text format. If this is
+	// set, it will be converted to a binary profile which will be subsequently used for dexpreopt.
+	rewrittenProfile android.Path
 }
 
 type DexpreoptProperties struct {
@@ -158,6 +164,11 @@ type DexpreoptProperties struct {
 		// defaults to searching for a file that matches the name of this module in the default
 		// profile location set by PRODUCT_DEX_PREOPT_PROFILE_DIR, or empty if not found.
 		Profile *string `android:"path"`
+
+		// If set to true, r8/d8 will use `profile` as input to generate a new profile that matches
+		// the optimized dex.
+		// The new profile will be subsequently used as the profile to dexpreopt the dex file.
+		Enable_profile_rewriting *bool
 	}
 
 	Dex_preopt_result struct {
@@ -421,13 +432,17 @@ func (d *dexpreopter) dexpreopt(ctx android.ModuleContext, libName string, dexJa
 	if d.inputProfilePathOnHost != nil {
 		profileClassListing = android.OptionalPathForPath(d.inputProfilePathOnHost)
 	} else if BoolDefault(d.dexpreoptProperties.Dex_preopt.Profile_guided, true) && !forPrebuiltApex(ctx) {
-		// If dex_preopt.profile_guided is not set, default it based on the existence of the
-		// dexprepot.profile option or the profile class listing.
-		if String(d.dexpreoptProperties.Dex_preopt.Profile) != "" {
+		// If enable_profile_rewriting is set, use the rewritten profile instead of the checked-in profile
+		if d.EnableProfileRewriting() {
+			profileClassListing = android.OptionalPathForPath(d.GetRewrittenProfile())
+			profileIsTextListing = true
+		} else if profile := d.GetProfile(); profile != "" {
+			// If dex_preopt.profile_guided is not set, default it based on the existence of the
+			// dexprepot.profile option or the profile class listing.
 			profileClassListing = android.OptionalPathForPath(
-				android.PathForModuleSrc(ctx, String(d.dexpreoptProperties.Dex_preopt.Profile)))
+				android.PathForModuleSrc(ctx, profile))
 			profileBootListing = android.ExistentPathForSource(ctx,
-				ctx.ModuleDir(), String(d.dexpreoptProperties.Dex_preopt.Profile)+"-boot")
+				ctx.ModuleDir(), profile+"-boot")
 			profileIsTextListing = true
 		} else if global.ProfileDir != "" {
 			profileClassListing = android.ExistentPathForSource(ctx,
@@ -588,3 +603,23 @@ func (d *dexpreopter) OutputProfilePathOnHost() android.Path {
 func (d *dexpreopter) disableDexpreopt() {
 	d.shouldDisableDexpreopt = true
 }
+
+func (d *dexpreopter) EnableProfileRewriting() bool {
+	return proptools.Bool(d.dexpreoptProperties.Dex_preopt.Enable_profile_rewriting)
+}
+
+func (d *dexpreopter) GetProfile() string {
+	return proptools.String(d.dexpreoptProperties.Dex_preopt.Profile)
+}
+
+func (d *dexpreopter) GetProfileGuided() bool {
+	return proptools.Bool(d.dexpreoptProperties.Dex_preopt.Profile_guided)
+}
+
+func (d *dexpreopter) GetRewrittenProfile() android.Path {
+	return d.rewrittenProfile
+}
+
+func (d *dexpreopter) SetRewrittenProfile(p android.Path) {
+	d.rewrittenProfile = p
+}

java/java.go

@@ -2339,7 +2339,7 @@ func (al *ApiLibrary) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 		classesJar:    al.stubsJar,
 		jarName:       ctx.ModuleName() + ".jar",
 	}
-	dexOutputFile := al.dexer.compileDex(ctx, dexParams)
+	dexOutputFile, _ := al.dexer.compileDex(ctx, dexParams)
 	uncompressed := true
 	al.initHiddenAPI(ctx, makeDexJarPathFromPath(dexOutputFile), al.stubsJar, &uncompressed)
 	dexOutputFile = al.hiddenAPIEncodeDex(ctx, dexOutputFile)
@@ -2723,7 +2723,7 @@ func (j *Import) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 				jarName:       jarName,
 			}
 
-			dexOutputFile = j.dexer.compileDex(ctx, dexParams)
+			dexOutputFile, _ = j.dexer.compileDex(ctx, dexParams)
 			if ctx.Failed() {
 				return
 			}