Use the correct bootjars when multiple prebuilt apexes exist

hiddenapi and dexpreopt require boot and system server jars from apexes.
When building with prebuilts, this comes via
java_import/java_sdk_library_import, which acts as a hook for
prebuilt_apex/apex_set. If we have multiple apexes in the tree, this
hook becomes 1:many. This CL prepares dex_bootjars to select the right
deapexerd .jar files when mutliple prebuilts exist.

Implementation details
- Update prebuilt module types (prebuilt_apex/apex_set) and source
  apexes to set a map of
  library name to dex jar path on host.
- dex_bootjars will access the path of the .dex jar on host via the
  provider. These then
  copied/installed to the right locations.

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

Bug: 308790457
Test: git_master-art-host:art-gtest https://android-build.corp.google.com/builds/abtd/run/L21500030000926533
Test: git_main:art_standalone_dexpreopt_tests https://android-build.corp.google.com/builds/abtd/run/L99000030000891212
Test: Added a unit test that checks that the right .jar is selected
when multiple prebuilts exists

Change-Id: I6ef94135b9303a35135810930af4b641df13a583
diff --git a/android/apex.go b/android/apex.go
index b4bb67c..c1e7a5c 100644
--- a/android/apex.go
+++ b/android/apex.go
@@ -965,4 +965,7 @@
 
 	// Path to the image profile file on host (or empty, if profile is not generated).
 	ProfilePathOnHost Path
+
+	// Map from the apex library name (without prebuilt_ prefix) to the dex file path on host
+	LibraryNameToDexJarPathOnHost map[string]Path
 }
diff --git a/android/deapexer.go b/android/deapexer.go
index fb2073d..2704b3e 100644
--- a/android/deapexer.go
+++ b/android/deapexer.go
@@ -79,6 +79,10 @@
 	//
 	// See Prebuilt.ApexInfoMutator for more information.
 	exports map[string]WritablePath
+
+	// name of the java libraries exported from the apex
+	// e.g. core-libart
+	exportedModuleNames []string
 }
 
 // ApexModuleName returns the name of the APEX module that provided the info.
@@ -97,6 +101,10 @@
 	return path
 }
 
+func (i DeapexerInfo) GetExportedModuleNames() []string {
+	return i.exportedModuleNames
+}
+
 // Provider that can be used from within the `GenerateAndroidBuildActions` of a module that depends
 // on a `deapexer` module to retrieve its `DeapexerInfo`.
 var DeapexerProvider = blueprint.NewProvider[DeapexerInfo]()
@@ -105,10 +113,11 @@
 // for use with a prebuilt_apex module.
 //
 // See apex/deapexer.go for more information.
-func NewDeapexerInfo(apexModuleName string, exports map[string]WritablePath) DeapexerInfo {
+func NewDeapexerInfo(apexModuleName string, exports map[string]WritablePath, moduleNames []string) DeapexerInfo {
 	return DeapexerInfo{
-		apexModuleName: apexModuleName,
-		exports:        exports,
+		apexModuleName:      apexModuleName,
+		exports:             exports,
+		exportedModuleNames: moduleNames,
 	}
 }
 
diff --git a/apex/apex.go b/apex/apex.go
index 35a8782..29d59e5 100644
--- a/apex/apex.go
+++ b/apex/apex.go
@@ -2382,8 +2382,9 @@
 	ctx.VisitDirectDepsWithTag(bcpfTag, func(child android.Module) {
 		if info, ok := android.OtherModuleProvider(ctx, child, java.BootclasspathFragmentApexContentInfoProvider); ok {
 			exports := android.ApexExportsInfo{
-				ApexName:          a.ApexVariationName(),
-				ProfilePathOnHost: info.ProfilePathOnHost(),
+				ApexName:                      a.ApexVariationName(),
+				ProfilePathOnHost:             info.ProfilePathOnHost(),
+				LibraryNameToDexJarPathOnHost: info.DexBootJarPathMap(),
 			}
 			ctx.SetProvider(android.ApexExportsInfoProvider, exports)
 		}
diff --git a/apex/apex_test.go b/apex/apex_test.go
index 6c7dabc..616421a 100644
--- a/apex/apex_test.go
+++ b/apex/apex_test.go
@@ -8408,30 +8408,32 @@
 func TestDuplicateDeapexersFromPrebuiltApexes(t *testing.T) {
 	preparers := android.GroupFixturePreparers(
 		java.PrepareForTestWithJavaDefaultModules,
+		prepareForTestWithBootclasspathFragment,
+		dexpreopt.FixtureSetTestOnlyArtBootImageJars("com.android.art:libfoo"),
 		PrepareForTestWithApexBuildComponents,
 	).
 		ExtendWithErrorHandler(android.FixtureExpectsAtLeastOneErrorMatchingPattern(
-			"Multiple installable prebuilt APEXes provide ambiguous deapexers: com.android.myapex and com.mycompany.android.myapex"))
+			"Multiple installable prebuilt APEXes provide ambiguous deapexers: com.android.art and com.mycompany.android.art"))
 
 	bpBase := `
 		apex_set {
-			name: "com.android.myapex",
+			name: "com.android.art",
 			installable: true,
-			exported_bootclasspath_fragments: ["my-bootclasspath-fragment"],
+			exported_bootclasspath_fragments: ["art-bootclasspath-fragment"],
 			set: "myapex.apks",
 		}
 
 		apex_set {
-			name: "com.mycompany.android.myapex",
-			apex_name: "com.android.myapex",
+			name: "com.mycompany.android.art",
+			apex_name: "com.android.art",
 			installable: true,
-			exported_bootclasspath_fragments: ["my-bootclasspath-fragment"],
+			exported_bootclasspath_fragments: ["art-bootclasspath-fragment"],
 			set: "company-myapex.apks",
 		}
 
 		prebuilt_bootclasspath_fragment {
-			name: "my-bootclasspath-fragment",
-			apex_available: ["com.android.myapex"],
+			name: "art-bootclasspath-fragment",
+			apex_available: ["com.android.art"],
 			hidden_api: {
 				annotation_flags: "my-bootclasspath-fragment/annotation-flags.csv",
 				metadata: "my-bootclasspath-fragment/metadata.csv",
@@ -8448,7 +8450,7 @@
 			java_import {
 				name: "libfoo",
 				jars: ["libfoo.jar"],
-				apex_available: ["com.android.myapex"],
+				apex_available: ["com.android.art"],
 			}
 		`)
 	})
@@ -8461,7 +8463,7 @@
 					jars: ["libbar.jar"],
 				},
 				shared_library: false,
-				apex_available: ["com.android.myapex"],
+				apex_available: ["com.android.art"],
 			}
 		`)
 	})
@@ -8477,7 +8479,7 @@
 					jars: ["libbar.jar"],
 				},
 				shared_library: false,
-				apex_available: ["com.android.myapex"],
+				apex_available: ["com.android.art"],
 			}
 		`)
 	})
@@ -11413,3 +11415,165 @@
 	android.EnsureListContainsSuffix(t, buildParams.Inputs.Strings(), "my_aconfig_declarations_foo/intermediate.pb")
 	ensureContains(t, buildParams.Output.String(), "android_common_myapex/aconfig_flags.pb")
 }
+
+// Test that the boot jars come from the _selected_ apex prebuilt
+// RELEASE_APEX_CONTIRBUTIONS_* build flags will be used to select the correct prebuilt for a specific release config
+func TestBootDexJarsMultipleApexPrebuilts(t *testing.T) {
+	checkBootDexJarPath := func(t *testing.T, ctx *android.TestContext, stem string, bootDexJarPath string) {
+		t.Helper()
+		s := ctx.ModuleForTests("dex_bootjars", "android_common")
+		foundLibfooJar := false
+		base := stem + ".jar"
+		for _, output := range s.AllOutputs() {
+			if filepath.Base(output) == base {
+				foundLibfooJar = true
+				buildRule := s.Output(output)
+				android.AssertStringEquals(t, "boot dex jar path", bootDexJarPath, buildRule.Input.String())
+			}
+		}
+		if !foundLibfooJar {
+			t.Errorf("Rule for libfoo.jar missing in dex_bootjars singleton outputs %q", android.StringPathsRelativeToTop(ctx.Config().SoongOutDir(), s.AllOutputs()))
+		}
+	}
+
+	bp := `
+		// Source APEX.
+
+		java_library {
+			name: "framework-foo",
+			srcs: ["foo.java"],
+			installable: true,
+			apex_available: [
+				"com.android.foo",
+			],
+		}
+
+		bootclasspath_fragment {
+			name: "foo-bootclasspath-fragment",
+			contents: ["framework-foo"],
+			apex_available: [
+				"com.android.foo",
+			],
+			hidden_api: {
+				split_packages: ["*"],
+			},
+		}
+
+		apex_key {
+			name: "com.android.foo.key",
+			public_key: "com.android.foo.avbpubkey",
+			private_key: "com.android.foo.pem",
+		}
+
+		apex {
+			name: "com.android.foo",
+			key: "com.android.foo.key",
+			bootclasspath_fragments: ["foo-bootclasspath-fragment"],
+			updatable: false,
+		}
+
+		// Prebuilt APEX.
+
+		java_sdk_library_import {
+			name: "framework-foo",
+			public: {
+				jars: ["foo.jar"],
+			},
+			apex_available: ["com.android.foo"],
+			shared_library: false,
+		}
+
+		prebuilt_bootclasspath_fragment {
+			name: "foo-bootclasspath-fragment",
+			contents: ["framework-foo"],
+			hidden_api: {
+				annotation_flags: "my-bootclasspath-fragment/annotation-flags.csv",
+				metadata: "my-bootclasspath-fragment/metadata.csv",
+				index: "my-bootclasspath-fragment/index.csv",
+				stub_flags: "my-bootclasspath-fragment/stub-flags.csv",
+				all_flags: "my-bootclasspath-fragment/all-flags.csv",
+			},
+			apex_available: [
+				"com.android.foo",
+			],
+		}
+
+		prebuilt_apex {
+			name: "com.android.foo",
+			apex_name: "com.android.foo",
+			src: "com.android.foo-arm.apex",
+			exported_bootclasspath_fragments: ["foo-bootclasspath-fragment"],
+		}
+
+		// Another Prebuilt ART APEX
+		prebuilt_apex {
+			name: "com.android.foo.v2",
+			apex_name: "com.android.foo", // Used to determine the API domain
+			src: "com.android.foo-arm.apex",
+			exported_bootclasspath_fragments: ["foo-bootclasspath-fragment"],
+		}
+
+		// APEX contribution modules
+
+		apex_contributions {
+			name: "foo.source.contributions",
+			api_domain: "com.android.foo",
+			contents: ["com.android.foo"],
+		}
+
+		apex_contributions {
+			name: "foo.prebuilt.contributions",
+			api_domain: "com.android.foo",
+			contents: ["prebuilt_com.android.foo"],
+		}
+
+		apex_contributions {
+			name: "foo.prebuilt.v2.contributions",
+			api_domain: "com.android.foo",
+			contents: ["com.android.foo.v2"], // prebuilt_ prefix is missing because of prebuilt_rename mutator
+		}
+	`
+
+	testCases := []struct {
+		desc                      string
+		selectedApexContributions string
+		expectedBootJar           string
+	}{
+		{
+			desc:                      "Source apex com.android.foo is selected, bootjar should come from source java library",
+			selectedApexContributions: "foo.source.contributions",
+			expectedBootJar:           "out/soong/.intermediates/foo-bootclasspath-fragment/android_common_apex10000/hiddenapi-modular/encoded/framework-foo.jar",
+		},
+		{
+			desc:                      "Prebuilt apex prebuilt_com.android.foo is selected, profile should come from .prof deapexed from the prebuilt",
+			selectedApexContributions: "foo.prebuilt.contributions",
+			expectedBootJar:           "out/soong/.intermediates/com.android.foo.deapexer/android_common/deapexer/javalib/framework-foo.jar",
+		},
+		{
+			desc:                      "Prebuilt apex prebuilt_com.android.foo.v2 is selected, profile should come from .prof deapexed from the prebuilt",
+			selectedApexContributions: "foo.prebuilt.v2.contributions",
+			expectedBootJar:           "out/soong/.intermediates/com.android.foo.v2.deapexer/android_common/deapexer/javalib/framework-foo.jar",
+		},
+	}
+
+	fragment := java.ApexVariantReference{
+		Apex:   proptools.StringPtr("com.android.foo"),
+		Module: proptools.StringPtr("foo-bootclasspath-fragment"),
+	}
+
+	for _, tc := range testCases {
+		preparer := android.GroupFixturePreparers(
+			java.FixtureConfigureApexBootJars("com.android.foo:framework-foo"),
+			android.FixtureMergeMockFs(map[string][]byte{
+				"system/sepolicy/apex/com.android.foo-file_contexts": nil,
+			}),
+			android.FixtureModifyProductVariables(func(variables android.FixtureProductVariables) {
+				variables.BuildFlags = map[string]string{
+					"RELEASE_APEX_CONTRIBUTIONS_ADSERVICES": tc.selectedApexContributions,
+				}
+			}),
+		)
+		ctx := testDexpreoptWithApexes(t, bp, "", preparer, fragment)
+		checkBootDexJarPath(t, ctx, "framework-foo", tc.expectedBootJar)
+	}
+}
diff --git a/apex/deapexer.go b/apex/deapexer.go
index 5aeea63..5ff622c 100644
--- a/apex/deapexer.go
+++ b/apex/deapexer.go
@@ -98,6 +98,7 @@
 func (p *Deapexer) DepsMutator(ctx android.BottomUpMutatorContext) {
 	// Add dependencies from the java modules to which this exports files from the `.apex` file onto
 	// this module so that they can access the `DeapexerInfo` object that this provides.
+	// TODO: b/308174306 - Once all the mainline modules have been flagged, drop this dependency edge
 	for _, lib := range p.properties.CommonModules {
 		dep := prebuiltApexExportedModuleName(ctx, lib)
 		ctx.AddReverseDependency(ctx.Module(), android.DeapexerTag, dep)
@@ -126,7 +127,7 @@
 	// apex relative path to extracted file path available for other modules.
 	if len(exports) > 0 {
 		// Make the information available for other modules.
-		di := android.NewDeapexerInfo(apexModuleName(ctx.ModuleName()), exports)
+		di := android.NewDeapexerInfo(apexModuleName(ctx.ModuleName()), exports, p.properties.CommonModules)
 		android.SetProvider(ctx, android.DeapexerProvider, di)
 
 		// Create a sorted list of the files that this exports.
diff --git a/apex/dexpreopt_bootjars_test.go b/apex/dexpreopt_bootjars_test.go
index 9d74519..34ccdd7 100644
--- a/apex/dexpreopt_bootjars_test.go
+++ b/apex/dexpreopt_bootjars_test.go
@@ -308,9 +308,18 @@
 
 		// Prebuilt ART APEX.
 
+		java_import {
+			name: "core-oj",
+			jars: ["core-oj.jar"],
+			apex_available: [
+				"com.android.art",
+			],
+		}
+
 		prebuilt_bootclasspath_fragment {
 			name: "art-bootclasspath-fragment",
 			image_name: "art",
+			contents: ["core-oj"],
 			hidden_api: {
 				annotation_flags: "my-bootclasspath-fragment/annotation-flags.csv",
 				metadata: "my-bootclasspath-fragment/metadata.csv",
diff --git a/apex/prebuilt.go b/apex/prebuilt.go
index 37a9ff5..1f57b63 100644
--- a/apex/prebuilt.go
+++ b/apex/prebuilt.go
@@ -778,9 +778,17 @@
 		return
 	}
 	if di, err := android.FindDeapexerProviderForModule(ctx); err == nil {
+		javaModuleToDexPath := map[string]android.Path{}
+		for _, commonModule := range di.GetExportedModuleNames() {
+			if dex := di.PrebuiltExportPath(java.ApexRootRelativePathToJavaLib(commonModule)); dex != nil {
+				javaModuleToDexPath[commonModule] = dex
+			}
+		}
+
 		exports := android.ApexExportsInfo{
-			ApexName:          p.ApexVariationName(),
-			ProfilePathOnHost: di.PrebuiltExportPath(java.ProfileInstallPathInApex),
+			ApexName:                      p.ApexVariationName(),
+			ProfilePathOnHost:             di.PrebuiltExportPath(java.ProfileInstallPathInApex),
+			LibraryNameToDexJarPathOnHost: javaModuleToDexPath,
 		}
 		ctx.SetProvider(android.ApexExportsInfoProvider, exports)
 	} else {
diff --git a/java/bootclasspath_fragment.go b/java/bootclasspath_fragment.go
index 010dbec..ae24404 100644
--- a/java/bootclasspath_fragment.go
+++ b/java/bootclasspath_fragment.go
@@ -385,6 +385,10 @@
 	}
 }
 
+func (i BootclasspathFragmentApexContentInfo) DexBootJarPathMap() bootDexJarByModule {
+	return i.contentModuleDexJarPaths
+}
+
 func (i BootclasspathFragmentApexContentInfo) ProfilePathOnHost() android.Path {
 	return i.profilePathOnHost
 }
@@ -1034,10 +1038,6 @@
 		return android.PathForModuleSrc(ctx, *src)
 	}
 
-	// Retrieve the dex files directly from the content modules. They in turn should retrieve the
-	// encoded dex jars from the prebuilt .apex files.
-	encodedBootDexJarsByModule := extractEncodedDexJarsFromModules(ctx, contents)
-
 	output := HiddenAPIOutput{
 		HiddenAPIFlagOutput: HiddenAPIFlagOutput{
 			AnnotationFlagsPath:   pathForSrc("hidden_api.annotation_flags", module.prebuiltProperties.Hidden_api.Annotation_flags),
@@ -1048,8 +1048,6 @@
 			StubFlagsPath: pathForOptionalSrc(module.prebuiltProperties.Hidden_api.Stub_flags, nil),
 			AllFlagsPath:  pathForOptionalSrc(module.prebuiltProperties.Hidden_api.All_flags, nil),
 		},
-
-		EncodedBootDexFilesByModule: encodedBootDexJarsByModule,
 	}
 
 	// TODO: Temporarily fallback to stub_flags/all_flags properties until prebuilts have been updated.
diff --git a/java/dexpreopt.go b/java/dexpreopt.go
index 1e289c5..bd3cce4 100644
--- a/java/dexpreopt.go
+++ b/java/dexpreopt.go
@@ -272,7 +272,7 @@
 	dc := dexpreopt.GetGlobalConfig(ctx)
 	d.installPath = android.PathForModuleInPartitionInstall(ctx, "", strings.TrimPrefix(dexpreopt.GetSystemServerDexLocation(ctx, dc, libraryName), "/"))
 	// generate the rules for creating the .odex and .vdex files for this system server jar
-	dexJarFile := di.PrebuiltExportPath(apexRootRelativePathToJavaLib(libraryName))
+	dexJarFile := di.PrebuiltExportPath(ApexRootRelativePathToJavaLib(libraryName))
 	d.dexpreopt(ctx, dexJarFile)
 }
 
diff --git a/java/dexpreopt_bootjars.go b/java/dexpreopt_bootjars.go
index e158ed3..205d3b9 100644
--- a/java/dexpreopt_bootjars.go
+++ b/java/dexpreopt_bootjars.go
@@ -699,38 +699,57 @@
 // extractEncodedDexJarsFromModulesOrBootclasspathFragments gets the hidden API encoded dex jars for
 // the given modules.
 func extractEncodedDexJarsFromModulesOrBootclasspathFragments(ctx android.ModuleContext, apexJarModulePairs []apexJarModulePair) bootDexJarByModule {
+	apexNameToBcpInfoMap := getApexNameToBcpInfoMap(ctx)
 	encodedDexJarsByModuleName := bootDexJarByModule{}
 	for _, pair := range apexJarModulePairs {
-		var path android.Path
-		if android.IsConfiguredJarForPlatform(pair.apex) || android.IsModulePrebuilt(pair.jarModule) {
-			// This gives us the dex jar with the hidden API flags encoded from the monolithic hidden API
-			// files or the dex jar extracted from a prebuilt APEX. We can't use this for a boot jar for
-			// a source APEX because there is no guarantee that it is the same as the jar packed into the
-			// APEX. In practice, they are the same when we are building from a full source tree, but they
-			// are different when we are building from a thin manifest (e.g., master-art), where there is
-			// no monolithic hidden API files at all.
-			path = retrieveEncodedBootDexJarFromModule(ctx, pair.jarModule)
-		} else {
-			// Use exactly the same jar that is packed into the APEX.
-			fragment := getBootclasspathFragmentByApex(ctx, pair.apex)
-			if fragment == nil {
-				ctx.ModuleErrorf("Boot jar '%[1]s' is from APEX '%[2]s', but a bootclasspath_fragment for "+
-					"APEX '%[2]s' doesn't exist or is not added as a dependency of dex_bootjars",
-					pair.jarModule.Name(),
-					pair.apex)
-			}
-			bootclasspathFragmentInfo, _ := android.OtherModuleProvider(ctx, fragment, BootclasspathFragmentApexContentInfoProvider)
-			jar, err := bootclasspathFragmentInfo.DexBootJarPathForContentModule(pair.jarModule)
-			if err != nil {
-				ctx.ModuleErrorf("%s", err)
-			}
-			path = jar
-		}
-		encodedDexJarsByModuleName.addPath(pair.jarModule, path)
+		dexJarPath := getDexJarForApex(ctx, pair, apexNameToBcpInfoMap)
+		encodedDexJarsByModuleName.addPath(pair.jarModule, dexJarPath)
 	}
 	return encodedDexJarsByModuleName
 }
 
+// Returns the java libraries exported by the apex for hiddenapi and dexpreopt
+// This information can come from two mechanisms
+// 1. New: Direct deps to _selected_ apexes. The apexes return a ApexExportsInfo
+// 2. Legacy: An edge to java_library or java_import (java_sdk_library) 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 getDexJarForApex(ctx android.ModuleContext, pair apexJarModulePair, apexNameToBcpInfoMap map[string]android.ApexExportsInfo) android.Path {
+	if info, exists := apexNameToBcpInfoMap[pair.apex]; exists {
+		libraryName := android.RemoveOptionalPrebuiltPrefix(pair.jarModule.Name())
+		if dex, exists := info.LibraryNameToDexJarPathOnHost[libraryName]; exists {
+			return dex
+		} else {
+			ctx.ModuleErrorf("Apex %s does not provide a dex boot jar for library %s\n", pair.apex, libraryName)
+		}
+	}
+	// TODO: b/308174306 - Remove the legacy mechanism
+	if android.IsConfiguredJarForPlatform(pair.apex) || android.IsModulePrebuilt(pair.jarModule) {
+		// This gives us the dex jar with the hidden API flags encoded from the monolithic hidden API
+		// files or the dex jar extracted from a prebuilt APEX. We can't use this for a boot jar for
+		// a source APEX because there is no guarantee that it is the same as the jar packed into the
+		// APEX. In practice, they are the same when we are building from a full source tree, but they
+		// are different when we are building from a thin manifest (e.g., master-art), where there is
+		// no monolithic hidden API files at all.
+		return retrieveEncodedBootDexJarFromModule(ctx, pair.jarModule)
+	} else {
+		// Use exactly the same jar that is packed into the APEX.
+		fragment := getBootclasspathFragmentByApex(ctx, pair.apex)
+		if fragment == nil {
+			ctx.ModuleErrorf("Boot jar '%[1]s' is from APEX '%[2]s', but a bootclasspath_fragment for "+
+				"APEX '%[2]s' doesn't exist or is not added as a dependency of dex_bootjars",
+				pair.jarModule.Name(),
+				pair.apex)
+		}
+		bootclasspathFragmentInfo, _ := android.OtherModuleProvider(ctx, fragment, BootclasspathFragmentApexContentInfoProvider)
+		jar, err := bootclasspathFragmentInfo.DexBootJarPathForContentModule(pair.jarModule)
+		if err != nil {
+			ctx.ModuleErrorf("%s", err)
+		}
+		return jar
+	}
+	return nil
+}
+
 // copyBootJarsToPredefinedLocations generates commands that will copy boot jars to predefined
 // paths in the global config.
 func copyBootJarsToPredefinedLocations(ctx android.ModuleContext, srcBootDexJarsByModule bootDexJarByModule, dstBootJarsByModule map[string]android.WritablePath) {
@@ -881,6 +900,16 @@
 	return fragment.(commonBootclasspathFragment).getProfilePath()
 }
 
+func getApexNameToBcpInfoMap(ctx android.ModuleContext) map[string]android.ApexExportsInfo {
+	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
+		}
+	})
+	return apexNameToBcpInfoMap
+}
+
 // Generate boot image build rules for a specific target.
 func buildBootImageVariant(ctx android.ModuleContext, image *bootImageVariant, profile android.Path) bootImageVariantOutputs {
 
@@ -923,12 +952,7 @@
 
 	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
-		}
-	})
+	apexNameToBcpInfoMap := getApexNameToBcpInfoMap(ctx)
 
 	cmd.Tool(globalSoong.Dex2oat).
 		Flag("--avoid-storing-invocation").
diff --git a/java/hiddenapi_modular.go b/java/hiddenapi_modular.go
index 8011f34..bf99757 100644
--- a/java/hiddenapi_modular.go
+++ b/java/hiddenapi_modular.go
@@ -947,6 +947,7 @@
 	HiddenAPIFlagOutput
 
 	// The map from base module name to the path to the encoded boot dex file.
+	// This field is not available in prebuilt apexes
 	EncodedBootDexFilesByModule bootDexJarByModule
 }
 
diff --git a/java/java.go b/java/java.go
index 51e8c34..2a4fafa 100644
--- a/java/java.go
+++ b/java/java.go
@@ -2258,11 +2258,11 @@
 				j.dexJarFileErr = err
 				return
 			}
-			dexJarFileApexRootRelative := apexRootRelativePathToJavaLib(j.BaseModuleName())
+			dexJarFileApexRootRelative := ApexRootRelativePathToJavaLib(j.BaseModuleName())
 			if dexOutputPath := di.PrebuiltExportPath(dexJarFileApexRootRelative); dexOutputPath != nil {
 				dexJarFile := makeDexJarPathFromPath(dexOutputPath)
 				j.dexJarFile = dexJarFile
-				installPath := android.PathForModuleInPartitionInstall(ctx, "apex", ai.ApexVariationName, apexRootRelativePathToJavaLib(j.BaseModuleName()))
+				installPath := android.PathForModuleInPartitionInstall(ctx, "apex", ai.ApexVariationName, ApexRootRelativePathToJavaLib(j.BaseModuleName()))
 				j.dexJarInstallFile = installPath
 
 				j.dexpreopter.installPath = j.dexpreopter.getInstallPath(ctx, installPath)
@@ -2422,7 +2422,7 @@
 // java_sdk_library_import with the specified base module name requires to be exported from a
 // prebuilt_apex/apex_set.
 func requiredFilesFromPrebuiltApexForImport(name string, d *dexpreopter) []string {
-	dexJarFileApexRootRelative := apexRootRelativePathToJavaLib(name)
+	dexJarFileApexRootRelative := ApexRootRelativePathToJavaLib(name)
 	// Add the dex implementation jar to the set of exported files.
 	files := []string{
 		dexJarFileApexRootRelative,
@@ -2433,9 +2433,9 @@
 	return files
 }
 
-// apexRootRelativePathToJavaLib returns the path, relative to the root of the apex's contents, for
+// ApexRootRelativePathToJavaLib returns the path, relative to the root of the apex's contents, for
 // the java library with the specified name.
-func apexRootRelativePathToJavaLib(name string) string {
+func ApexRootRelativePathToJavaLib(name string) string {
 	return filepath.Join("javalib", name+".jar")
 }
 
diff --git a/java/sdk_library.go b/java/sdk_library.go
index 38bd301..ef34fb6 100644
--- a/java/sdk_library.go
+++ b/java/sdk_library.go
@@ -2695,7 +2695,7 @@
 				module.dexJarFileErr = err
 				return
 			}
-			dexJarFileApexRootRelative := apexRootRelativePathToJavaLib(module.BaseModuleName())
+			dexJarFileApexRootRelative := ApexRootRelativePathToJavaLib(module.BaseModuleName())
 			if dexOutputPath := di.PrebuiltExportPath(dexJarFileApexRootRelative); dexOutputPath != nil {
 				dexJarFile := makeDexJarPathFromPath(dexOutputPath)
 				module.dexJarFile = dexJarFile