Use the correct prof file when multiple prebuilt apexes exist

Generating boot image requires a .prof file provided by the ART apex.
When building with prebuilts, this comes via the
prebuilt_bootclasspath_fragment module, which acts as a shim for
prebuilt_apex/apex_set. If we have multiple prebuilt apexes in the tree,
this shim becomes 1:many. This CL prepares dex_bootjars to select the
right .prof file when multiple prebuilts exist.

Implementation details
- Update deps mutator of dex_bootjars to create a dep on
  all_apex_contributions. The latter contains information about which
  apex is selected in a specific release configuration. dex_bootjars
  will create a dependency on the selected apex in a postdeps phase
  mutator.
- All apex module types (apex, prebuilt_apex and apex_set) will set a
  provider that contains info about the location of the .prof file on
  host
- dex_bootjars will access the provider of the selected apex to get the
  location of the .prof file

This CL does not drop the old mechanism to get the .prof file (i.e. by
creating a dep on {prebuilt_}bootclasspath_fragment). Once all mainline
modules have been flagged using apex_contributions, the old mechanism
will be dropped

Bug: 308790457
Test: Added a unit test that checks that the right .prof is selected
when multiple prebuilts exists

Change-Id: I40fdb21416c46bed32f6ff187ce5153711ec2c69
diff --git a/java/bootclasspath_fragment.go b/java/bootclasspath_fragment.go
index 83030b5..010dbec 100644
--- a/java/bootclasspath_fragment.go
+++ b/java/bootclasspath_fragment.go
@@ -534,7 +534,7 @@
 
 	if profile != nil {
 		info.profilePathOnHost = profile
-		info.profileInstallPathInApex = profileInstallPathInApex
+		info.profileInstallPathInApex = ProfileInstallPathInApex
 	}
 
 	// Make the apex content info available for other modules.
@@ -1074,7 +1074,7 @@
 		return nil // An error has been reported by FindDeapexerProviderForModule.
 	}
 
-	return di.PrebuiltExportPath(profileInstallPathInApex)
+	return di.PrebuiltExportPath(ProfileInstallPathInApex)
 }
 
 func (b *PrebuiltBootclasspathFragmentModule) getProfilePath() android.Path {
@@ -1094,7 +1094,7 @@
 func (module *PrebuiltBootclasspathFragmentModule) RequiredFilesFromPrebuiltApex(ctx android.BaseModuleContext) []string {
 	for _, apex := range module.ApexProperties.Apex_available {
 		if isProfileProviderApex(ctx, apex) {
-			return []string{profileInstallPathInApex}
+			return []string{ProfileInstallPathInApex}
 		}
 	}
 	return nil
diff --git a/java/dexpreopt_bootjars.go b/java/dexpreopt_bootjars.go
index 5a19945..e158ed3 100644
--- a/java/dexpreopt_bootjars.go
+++ b/java/dexpreopt_bootjars.go
@@ -21,6 +21,7 @@
 	"android/soong/android"
 	"android/soong/dexpreopt"
 
+	"github.com/google/blueprint"
 	"github.com/google/blueprint/proptools"
 )
 
@@ -224,8 +225,9 @@
 }
 
 var (
-	dexpreoptBootJarDepTag  = bootclasspathDependencyTag{name: "dexpreopt-boot-jar"}
-	dexBootJarsFragmentsKey = android.NewOnceKey("dexBootJarsFragments")
+	dexpreoptBootJarDepTag          = bootclasspathDependencyTag{name: "dexpreopt-boot-jar"}
+	dexBootJarsFragmentsKey         = android.NewOnceKey("dexBootJarsFragments")
+	apexContributionsMetadataDepTag = dependencyTag{name: "all_apex_contributions"}
 )
 
 func init() {
@@ -502,6 +504,11 @@
 	dexpreoptConfigForMake android.WritablePath
 }
 
+func (dbj *dexpreoptBootJars) DepsMutator(ctx android.BottomUpMutatorContext) {
+	// Create a dependency on all_apex_contributions to determine the selected mainline module
+	ctx.AddDependency(ctx.Module(), apexContributionsMetadataDepTag, "all_apex_contributions")
+}
+
 func DexpreoptBootJarsMutator(ctx android.BottomUpMutatorContext) {
 	if _, ok := ctx.Module().(*dexpreoptBootJars); !ok {
 		return
@@ -520,6 +527,14 @@
 		}
 		// For accessing the boot jars.
 		addDependenciesOntoBootImageModules(ctx, config.modules, dexpreoptBootJarDepTag)
+		// Create a dependency on the apex selected using RELEASE_APEX_CONTRIBUTIONS_*
+		// TODO: b/308174306 - Remove the direct depedendency edge to the java_library (source/prebuilt) once all mainline modules
+		// have been flagged using RELEASE_APEX_CONTRIBUTIONS_*
+		apexes := []string{}
+		for i := 0; i < config.modules.Len(); i++ {
+			apexes = append(apexes, config.modules.Apex(i))
+		}
+		addDependenciesOntoSelectedBootImageApexes(ctx, android.FirstUniqueStrings(apexes)...)
 	}
 
 	if ctx.OtherModuleExists("platform-bootclasspath") {
@@ -532,6 +547,28 @@
 	}
 }
 
+// Create a dependency from dex_bootjars to the specific apexes selected using all_apex_contributions
+// This dependency will be used to get the path to the deapexed dex boot jars and profile (via a provider)
+func addDependenciesOntoSelectedBootImageApexes(ctx android.BottomUpMutatorContext, apexes ...string) {
+	psi := android.PrebuiltSelectionInfoMap{}
+	ctx.VisitDirectDepsWithTag(apexContributionsMetadataDepTag, func(am android.Module) {
+		if ctx.OtherModuleHasProvider(am, android.PrebuiltSelectionInfoProvider) {
+			psi = ctx.OtherModuleProvider(am, android.PrebuiltSelectionInfoProvider).(android.PrebuiltSelectionInfoMap)
+		}
+	})
+	for _, apex := range apexes {
+		for _, selected := range psi.GetSelectedModulesForApiDomain(apex) {
+			// We need to add a dep on only the apex listed in `contents` of the selected apex_contributions module
+			// This is not available in a structured format in `apex_contributions`, so this hack adds a dep on all `contents`
+			// (some modules like art.module.public.api do not have an apex variation since it is a pure stub module that does not get installed)
+			apexVariationOfSelected := append(ctx.Target().Variations(), blueprint.Variation{Mutator: "apex", Variation: apex})
+			if ctx.OtherModuleDependencyVariantExists(apexVariationOfSelected, selected) {
+				ctx.AddFarVariationDependencies(apexVariationOfSelected, dexpreoptBootJarDepTag, selected)
+			}
+		}
+	}
+}
+
 func gatherBootclasspathFragments(ctx android.ModuleContext) map[string]android.Module {
 	return ctx.Config().Once(dexBootJarsFragmentsKey, func() interface{} {
 		fragments := make(map[string]android.Module)
@@ -823,6 +860,27 @@
 	config *bootImageVariant
 }
 
+// Returns the profile file for an apex
+// This information can come from two mechanisms
+// 1. New: Direct deps to _selected_ apexes. The apexes return a BootclasspathFragmentApexContentInfo
+// 2. Legacy: An edge to bootclasspath_fragment module. For prebuilt apexes, this serves as a hook and is populated by deapexers of prebuilt apxes
+// TODO: b/308174306 - Once all mainline modules have been flagged, drop (2)
+func getProfilePathForApex(ctx android.ModuleContext, apexName string, apexNameToBcpInfoMap map[string]android.ApexExportsInfo) android.Path {
+	if info, exists := apexNameToBcpInfoMap[apexName]; exists {
+		return info.ProfilePathOnHost
+	}
+	// TODO: b/308174306 - Remove the legacy mechanism
+	fragment := getBootclasspathFragmentByApex(ctx, apexName)
+	if fragment == nil {
+		ctx.ModuleErrorf("Boot image config imports profile from '%[2]s', but a "+
+			"bootclasspath_fragment for APEX '%[2]s' doesn't exist or is not added as a "+
+			"dependency of dex_bootjars",
+			apexName)
+		return nil
+	}
+	return fragment.(commonBootclasspathFragment).getProfilePath()
+}
+
 // Generate boot image build rules for a specific target.
 func buildBootImageVariant(ctx android.ModuleContext, image *bootImageVariant, profile android.Path) bootImageVariantOutputs {
 
@@ -865,6 +923,13 @@
 
 	invocationPath := outputPath.ReplaceExtension(ctx, "invocation")
 
+	apexNameToBcpInfoMap := map[string]android.ApexExportsInfo{}
+	ctx.VisitDirectDepsWithTag(dexpreoptBootJarDepTag, func(am android.Module) {
+		if info, exists := android.OtherModuleProvider(ctx, am, android.ApexExportsInfoProvider); exists {
+			apexNameToBcpInfoMap[info.ApexName] = info
+		}
+	})
+
 	cmd.Tool(globalSoong.Dex2oat).
 		Flag("--avoid-storing-invocation").
 		FlagWithOutput("--write-invocation-to=", invocationPath).ImplicitOutput(invocationPath).
@@ -877,16 +942,7 @@
 		}
 
 		for _, apex := range image.profileImports {
-			fragment := getBootclasspathFragmentByApex(ctx, apex)
-			if fragment == nil {
-				ctx.ModuleErrorf("Boot image config '%[1]s' imports profile from '%[2]s', but a "+
-					"bootclasspath_fragment for APEX '%[2]s' doesn't exist or is not added as a "+
-					"dependency of dex_bootjars",
-					image.name,
-					apex)
-				return bootImageVariantOutputs{}
-			}
-			importedProfile := fragment.(commonBootclasspathFragment).getProfilePath()
+			importedProfile := getProfilePathForApex(ctx, apex, apexNameToBcpInfoMap)
 			if importedProfile == nil {
 				ctx.ModuleErrorf("Boot image config '%[1]s' imports profile from '%[2]s', but '%[2]s' "+
 					"doesn't provide a profile",
diff --git a/java/dexpreopt_config.go b/java/dexpreopt_config.go
index 2bd696c..254b2c1 100644
--- a/java/dexpreopt_config.go
+++ b/java/dexpreopt_config.go
@@ -45,7 +45,7 @@
 	frameworkBootImageName   = "boot"
 	mainlineBootImageName    = "mainline"
 	bootImageStem            = "boot"
-	profileInstallPathInApex = "etc/boot-image.prof"
+	ProfileInstallPathInApex = "etc/boot-image.prof"
 )
 
 // getImageNames returns an ordered list of image names. The order doesn't matter but needs to be
diff --git a/java/sdk_library_test.go b/java/sdk_library_test.go
index 0965fc2..63419d6 100644
--- a/java/sdk_library_test.go
+++ b/java/sdk_library_test.go
@@ -950,6 +950,7 @@
 	})
 
 	CheckModuleDependencies(t, result.TestContext, "prebuilt_sdklib", "android_common", []string{
+		`all_apex_contributions`,
 		`prebuilt_sdklib.stubs`,
 		`sdklib.impl`,
 		// This should be prebuilt_sdklib.stubs but is set to sdklib.stubs because the
@@ -1022,6 +1023,7 @@
 	})
 
 	CheckModuleDependencies(t, result.TestContext, "prebuilt_sdklib", "android_common", []string{
+		`all_apex_contributions`,
 		`dex2oatd`,
 		`prebuilt_sdklib.stubs`,
 		`prebuilt_sdklib.stubs.source`,
@@ -1085,9 +1087,6 @@
 				"prebuilt_sdklib.source_preferred_using_legacy_flags",
 			],
 		}
-		all_apex_contributions {
-			name: "all_apex_contributions",
-		}
 		java_sdk_library {
 			name: "sdklib.prebuilt_preferred_using_legacy_flags",
 			srcs: ["a.java"],
@@ -1169,9 +1168,6 @@
 		prepareForJavaTest,
 		PrepareForTestWithJavaSdkLibraryFiles,
 		FixtureWithLastReleaseApis("sdklib.source_preferred_using_legacy_flags", "sdklib.prebuilt_preferred_using_legacy_flags"),
-		android.FixtureRegisterWithContext(func(ctx android.RegistrationContext) {
-			android.RegisterApexContributionsBuildComponents(ctx)
-		}),
 		android.FixtureModifyProductVariables(func(variables android.FixtureProductVariables) {
 			variables.BuildFlags = map[string]string{
 				"RELEASE_APEX_CONTRIBUTIONS_ADSERVICES": "my_mainline_module_contributions",
diff --git a/java/testing.go b/java/testing.go
index d55cffc..5959c49 100644
--- a/java/testing.go
+++ b/java/testing.go
@@ -383,6 +383,7 @@
 	RegisterSystemModulesBuildComponents(ctx)
 	registerSystemserverClasspathBuildComponents(ctx)
 	registerLintBuildComponents(ctx)
+	android.RegisterApexContributionsBuildComponents(ctx)
 }
 
 // gatherRequiredDepsForTest gathers the module definitions used by
@@ -570,6 +571,11 @@
 		}
 `
 
+	bp += `
+		all_apex_contributions {
+			name: "all_apex_contributions",
+		}
+`
 	return bp
 }