Pass libraryToApex and apexNameToFragment mappings into CreateClasspathElements

Remove a usage of ApexInfo.InApexVariants by collecting the libraryToApex
and apexToFragment mappings while collecting the fragments and passing
them into CreateClasspathElements.

Bug: 372543712
Test: CreateClasspathElementTest
Change-Id: I03adc821b04bc01828f075f25bbb8124505859a7
diff --git a/apex/classpath_element_test.go b/apex/classpath_element_test.go
index 55f1475..c2f2fc5 100644
--- a/apex/classpath_element_test.go
+++ b/apex/classpath_element_test.go
@@ -205,8 +205,6 @@
 	myFragment := result.Module("mybootclasspath-fragment", "android_common_myapex")
 	myBar := result.Module("bar", "android_common_apex10000")
 
-	other := result.Module("othersdklibrary", "android_common_apex10000")
-
 	otherApexLibrary := result.Module("otherapexlibrary", "android_common_apex10000")
 
 	platformFoo := result.Module("quuz", "android_common")
@@ -240,7 +238,11 @@
 	t.Run("art:baz, art:quuz, my:bar, foo", func(t *testing.T) {
 		t.Parallel()
 		ctx := newCtx()
-		elements := java.CreateClasspathElements(ctx, []android.Module{artBaz, artQuuz, myBar, platformFoo}, []android.Module{artFragment, myFragment})
+		elements := java.CreateClasspathElements(ctx,
+			[]android.Module{artBaz, artQuuz, myBar, platformFoo},
+			[]android.Module{artFragment, myFragment},
+			map[android.Module]string{artBaz: "com.android.art", artQuuz: "com.android.art", myBar: "myapex"},
+			map[string]android.Module{"com.android.art": artFragment, "myapex": myFragment})
 		expectedElements := java.ClasspathElements{
 			expectFragmentElement(artFragment, artBaz, artQuuz),
 			expectFragmentElement(myFragment, myBar),
@@ -249,32 +251,16 @@
 		assertElementsEquals(t, "elements", expectedElements, elements)
 	})
 
-	// Verify that CreateClasspathElements detects when an apex has multiple fragments.
-	t.Run("multiple fragments for same apex", func(t *testing.T) {
-		t.Parallel()
-		ctx := newCtx()
-		elements := java.CreateClasspathElements(ctx, []android.Module{}, []android.Module{artFragment, artFragment})
-		android.FailIfNoMatchingErrors(t, "apex com.android.art has multiple fragments, art-bootclasspath-fragment{.*} and art-bootclasspath-fragment{.*}", ctx.errs)
-		expectedElements := java.ClasspathElements{}
-		assertElementsEquals(t, "elements", expectedElements, elements)
-	})
-
-	// Verify that CreateClasspathElements detects when a library is in multiple fragments.
-	t.Run("library from multiple fragments", func(t *testing.T) {
-		t.Parallel()
-		ctx := newCtx()
-		elements := java.CreateClasspathElements(ctx, []android.Module{other}, []android.Module{artFragment, myFragment})
-		android.FailIfNoMatchingErrors(t, "library othersdklibrary{.*} is in two separate fragments, art-bootclasspath-fragment{.*} and mybootclasspath-fragment{.*}", ctx.errs)
-		expectedElements := java.ClasspathElements{}
-		assertElementsEquals(t, "elements", expectedElements, elements)
-	})
-
 	// Verify that CreateClasspathElements detects when a fragment's contents are not contiguous and
 	// are separated by a library from another fragment.
 	t.Run("discontiguous separated by fragment", func(t *testing.T) {
 		t.Parallel()
 		ctx := newCtx()
-		elements := java.CreateClasspathElements(ctx, []android.Module{artBaz, myBar, artQuuz, platformFoo}, []android.Module{artFragment, myFragment})
+		elements := java.CreateClasspathElements(ctx,
+			[]android.Module{artBaz, myBar, artQuuz, platformFoo},
+			[]android.Module{artFragment, myFragment},
+			map[android.Module]string{artBaz: "com.android.art", artQuuz: "com.android.art", myBar: "myapex"},
+			map[string]android.Module{"com.android.art": artFragment, "myapex": myFragment})
 		expectedElements := java.ClasspathElements{
 			expectFragmentElement(artFragment, artBaz, artQuuz),
 			expectFragmentElement(myFragment, myBar),
@@ -289,7 +275,11 @@
 	t.Run("discontiguous separated by library", func(t *testing.T) {
 		t.Parallel()
 		ctx := newCtx()
-		elements := java.CreateClasspathElements(ctx, []android.Module{artBaz, platformFoo, artQuuz, myBar}, []android.Module{artFragment, myFragment})
+		elements := java.CreateClasspathElements(ctx,
+			[]android.Module{artBaz, platformFoo, artQuuz, myBar},
+			[]android.Module{artFragment, myFragment},
+			map[android.Module]string{artBaz: "com.android.art", artQuuz: "com.android.art", myBar: "myapex"},
+			map[string]android.Module{"com.android.art": artFragment, "myapex": myFragment})
 		expectedElements := java.ClasspathElements{
 			expectFragmentElement(artFragment, artBaz, artQuuz),
 			expectLibraryElement(platformFoo),
@@ -305,7 +295,11 @@
 	t.Run("no fragment for apex", func(t *testing.T) {
 		t.Parallel()
 		ctx := newCtx()
-		elements := java.CreateClasspathElements(ctx, []android.Module{artBaz, otherApexLibrary}, []android.Module{artFragment})
+		elements := java.CreateClasspathElements(ctx,
+			[]android.Module{artBaz, otherApexLibrary},
+			[]android.Module{artFragment},
+			map[android.Module]string{artBaz: "com.android.art", otherApexLibrary: "otherapex"},
+			map[string]android.Module{"com.android.art": artFragment})
 		expectedElements := java.ClasspathElements{
 			expectFragmentElement(artFragment, artBaz),
 		}
diff --git a/java/bootclasspath.go b/java/bootclasspath.go
index 5500926..98fb417 100644
--- a/java/bootclasspath.go
+++ b/java/bootclasspath.go
@@ -68,10 +68,78 @@
 
 }
 
+// gatherFragments collects fragments that are direct dependencies of this module, as well as
+// any fragments in apexes via the dependency on the apex.  It returns a list of the fragment
+// modules and map from apex name to the fragment in that apex.
+func gatherFragments(ctx android.BaseModuleContext) ([]android.Module, map[string]android.Module) {
+	var fragments []android.Module
+
+	type fragmentInApex struct {
+		module string
+		apex   string
+	}
+
+	var fragmentsInApexes []fragmentInApex
+
+	// Find any direct dependencies, as well as a list of the modules in apexes.
+	ctx.VisitDirectDeps(func(module android.Module) {
+		t := ctx.OtherModuleDependencyTag(module)
+		if bcpTag, ok := t.(bootclasspathDependencyTag); ok && bcpTag.typ == fragment {
+			if bcpTag.moduleInApex != "" {
+				fragmentsInApexes = append(fragmentsInApexes, fragmentInApex{bcpTag.moduleInApex, ctx.OtherModuleName(module)})
+			} else {
+				fragments = append(fragments, module)
+			}
+		}
+	})
+
+	fragmentsMap := make(map[string]android.Module)
+	for _, fragmentInApex := range fragmentsInApexes {
+		var found android.Module
+		// Find a desired module in an apex.
+		ctx.WalkDeps(func(child, parent android.Module) bool {
+			t := ctx.OtherModuleDependencyTag(child)
+			if parent == ctx.Module() {
+				if bcpTag, ok := t.(bootclasspathDependencyTag); ok && bcpTag.typ == fragment && ctx.OtherModuleName(child) == fragmentInApex.apex {
+					// This is the dependency from this module to the apex, recurse into it.
+					return true
+				}
+			} else if android.IsDontReplaceSourceWithPrebuiltTag(t) {
+				return false
+			} else if t == android.PrebuiltDepTag {
+				return false
+			} else if IsBootclasspathFragmentContentDepTag(t) {
+				return false
+			} else if android.RemoveOptionalPrebuiltPrefix(ctx.OtherModuleName(child)) == fragmentInApex.module {
+				// This is the desired module inside the apex.
+				if found != nil && child != found {
+					panic(fmt.Errorf("found two conflicting modules %q in apex %q: %s and %s",
+						fragmentInApex.module, fragmentInApex.apex, found, child))
+				}
+				found = child
+			}
+			return false
+		})
+		if found != nil {
+			if existing, exists := fragmentsMap[fragmentInApex.apex]; exists {
+				ctx.ModuleErrorf("apex %s has multiple fragments, %s and %s", fragmentInApex.apex, fragmentInApex.module, existing)
+			} else {
+				fragmentsMap[fragmentInApex.apex] = found
+				fragments = append(fragments, found)
+			}
+		} else if !ctx.Config().AllowMissingDependencies() {
+			ctx.ModuleErrorf("failed to find fragment %q in apex %q\n",
+				fragmentInApex.module, fragmentInApex.apex)
+		}
+	}
+	return fragments, fragmentsMap
+}
+
 // gatherApexModulePairDepsWithTag returns the list of dependencies with the supplied tag that was
 // added by addDependencyOntoApexModulePair.
-func gatherApexModulePairDepsWithTag(ctx android.BaseModuleContext, tagType bootclasspathDependencyTagType) []android.Module {
+func gatherApexModulePairDepsWithTag(ctx android.BaseModuleContext, tagType bootclasspathDependencyTagType) ([]android.Module, map[android.Module]string) {
 	var modules []android.Module
+	modulesToApex := make(map[android.Module]string)
 
 	type moduleInApex struct {
 		module string
@@ -119,12 +187,17 @@
 		})
 		if found != nil {
 			modules = append(modules, found)
+			if existing, exists := modulesToApex[found]; exists && existing != moduleInApex.apex {
+				ctx.ModuleErrorf("module %s is in two apexes, %s and %s", moduleInApex.module, existing, moduleInApex.apex)
+			} else {
+				modulesToApex[found] = moduleInApex.apex
+			}
 		} else if !ctx.Config().AllowMissingDependencies() {
 			ctx.ModuleErrorf("failed to find module %q in apex %q\n",
 				moduleInApex.module, moduleInApex.apex)
 		}
 	}
-	return modules
+	return modules, modulesToApex
 }
 
 // ApexVariantReference specifies a particular apex variant of a module.
diff --git a/java/bootclasspath_fragment.go b/java/bootclasspath_fragment.go
index f3bff12..fd4b0a3 100644
--- a/java/bootclasspath_fragment.go
+++ b/java/bootclasspath_fragment.go
@@ -513,7 +513,7 @@
 		}
 	})
 
-	fragments := gatherApexModulePairDepsWithTag(ctx, fragment)
+	fragments, _ := gatherFragments(ctx)
 
 	// Perform hidden API processing.
 	hiddenAPIOutput := b.generateHiddenAPIBuildActions(ctx, contents, fragments)
diff --git a/java/classpath_element.go b/java/classpath_element.go
index abbcae7..4af2770 100644
--- a/java/classpath_element.go
+++ b/java/classpath_element.go
@@ -108,33 +108,18 @@
 //
 // e.g. Given the following input:
 //
-//	libraries: com.android.art:core-oj, com.android.art:core-libart, framework, ext
-//	fragments: com.android.art:art-bootclasspath-fragment
+//		libraries: core-oj, core-libart, framework, ext
+//		fragments: art-bootclasspath-fragment
+//	    libraryToApex: core-oj: com.android.art, core-libart: com.android.art
+//	    apexNameToFragment: com.android.art: art-bootclasspath-fragment
 //
 // Then this will return:
 //
 //	ClasspathFragmentElement(art-bootclasspath-fragment, [core-oj, core-libart]),
 //	ClasspathLibraryElement(framework),
 //	ClasspathLibraryElement(ext),
-func CreateClasspathElements(ctx ClasspathElementContext, libraries []android.Module, fragments []android.Module) ClasspathElements {
-	// Create a map from apex name to the fragment module. This makes it easy to find the fragment
-	// associated with a particular apex.
-	apexToFragment := map[string]android.Module{}
-	for _, fragment := range fragments {
-		apexInfo, ok := android.OtherModuleProvider(ctx, fragment, android.ApexInfoProvider)
-		if !ok {
-			ctx.ModuleErrorf("fragment %s is not part of an apex", fragment)
-			continue
-		}
-
-		for _, apex := range apexInfo.InApexVariants {
-			if existing, ok := apexToFragment[apex]; ok {
-				ctx.ModuleErrorf("apex %s has multiple fragments, %s and %s", apex, fragment, existing)
-				continue
-			}
-			apexToFragment[apex] = fragment
-		}
-	}
+func CreateClasspathElements(ctx ClasspathElementContext, libraries []android.Module, fragments []android.Module,
+	libraryToApex map[android.Module]string, apexNameToFragment map[string]android.Module) ClasspathElements {
 
 	fragmentToElement := map[android.Module]*ClasspathFragmentElement{}
 	elements := []ClasspathElement{}
@@ -144,31 +129,28 @@
 	// Iterate over the libraries to construct the ClasspathElements list.
 	for _, library := range libraries {
 		var element ClasspathElement
-		if apexInfo, ok := android.OtherModuleProvider(ctx, library, android.ApexInfoProvider); ok {
-
+		if libraryApex, ok := libraryToApex[library]; ok {
 			var fragment android.Module
 
 			// Make sure that the library is in only one fragment of the classpath.
-			for _, apex := range apexInfo.InApexVariants {
-				if f, ok := apexToFragment[apex]; ok {
-					if fragment == nil {
-						// This is the first fragment so just save it away.
-						fragment = f
-					} else if f != fragment {
-						// This apex variant of the library is in a different fragment.
-						ctx.ModuleErrorf("library %s is in two separate fragments, %s and %s", library, fragment, f)
-						// Skip over this library entirely as otherwise the resulting classpath elements would
-						// be invalid.
-						continue skipLibrary
-					}
-				} else {
-					// There is no fragment associated with the library's apex.
+			if f, ok := apexNameToFragment[libraryApex]; ok {
+				if fragment == nil {
+					// This is the first fragment so just save it away.
+					fragment = f
+				} else if f != fragment {
+					// This apex variant of the library is in a different fragment.
+					ctx.ModuleErrorf("library %s is in two separate fragments, %s and %s", library, fragment, f)
+					// Skip over this library entirely as otherwise the resulting classpath elements would
+					// be invalid.
+					continue skipLibrary
 				}
+			} else {
+				// There is no fragment associated with the library's apex.
 			}
 
 			if fragment == nil {
 				ctx.ModuleErrorf("library %s is from apexes %s which have no corresponding fragment in %s",
-					library, apexInfo.InApexVariants, fragments)
+					library, []string{libraryApex}, fragments)
 				// Skip over this library entirely as otherwise the resulting classpath elements would
 				// be invalid.
 				continue skipLibrary
diff --git a/java/dexpreopt_bootjars.go b/java/dexpreopt_bootjars.go
index 24bb99d..27027f0 100644
--- a/java/dexpreopt_bootjars.go
+++ b/java/dexpreopt_bootjars.go
@@ -734,7 +734,8 @@
 	modules := make([]apexJarModulePair, 0, imageConfig.modules.Len())
 	for i := 0; i < imageConfig.modules.Len(); i++ {
 		found := false
-		for _, module := range gatherApexModulePairDepsWithTag(ctx, dexpreoptBootJar) {
+		dexpreoptBootJarModules, _ := gatherApexModulePairDepsWithTag(ctx, dexpreoptBootJar)
+		for _, module := range dexpreoptBootJarModules {
 			name := android.RemoveOptionalPrebuiltPrefix(module.Name())
 			if name == imageConfig.modules.Jar(i) {
 				modules = append(modules, apexJarModulePair{
@@ -817,6 +818,7 @@
 				"APEX '%[2]s' doesn't exist or is not added as a dependency of dex_bootjars",
 				pair.jarModule.Name(),
 				pair.apex)
+			return nil
 		}
 		bootclasspathFragmentInfo, _ := android.OtherModuleProvider(ctx, fragment, BootclasspathFragmentApexContentInfoProvider)
 		jar, err := bootclasspathFragmentInfo.DexBootJarPathForContentModule(pair.jarModule)
diff --git a/java/platform_bootclasspath.go b/java/platform_bootclasspath.go
index 3b05b16..39b54e3 100644
--- a/java/platform_bootclasspath.go
+++ b/java/platform_bootclasspath.go
@@ -15,6 +15,9 @@
 package java
 
 import (
+	"maps"
+	"slices"
+
 	"github.com/google/blueprint"
 
 	"android/soong/android"
@@ -53,6 +56,12 @@
 	// The apex:module pairs obtained from the fragments.
 	fragments []android.Module
 
+	// The map of apex to the fragments they contain.
+	apexNameToFragment map[string]android.Module
+
+	// The map of library modules in the bootclasspath to the fragments that contain them.
+	libraryToApex map[android.Module]string
+
 	// Path to the monolithic hiddenapi-flags.csv file.
 	hiddenAPIFlagsCSV android.OutputPath
 
@@ -159,16 +168,16 @@
 
 func (b *platformBootclasspathModule) GenerateAndroidBuildActions(ctx android.ModuleContext) {
 	// Gather all the dependencies from the art, platform, and apex boot jars.
-	artModules := gatherApexModulePairDepsWithTag(ctx, artBootJar)
-	platformModules := gatherApexModulePairDepsWithTag(ctx, platformBootJar)
-	apexModules := gatherApexModulePairDepsWithTag(ctx, apexBootJar)
+	artModules, artModulesToApex := gatherApexModulePairDepsWithTag(ctx, artBootJar)
+	platformModules, platformModulesToApex := gatherApexModulePairDepsWithTag(ctx, platformBootJar)
+	apexModules, apexModulesToApex := gatherApexModulePairDepsWithTag(ctx, apexBootJar)
 
 	// Concatenate them all, in order as they would appear on the bootclasspath.
-	var allModules []android.Module
-	allModules = append(allModules, artModules...)
-	allModules = append(allModules, platformModules...)
-	allModules = append(allModules, apexModules...)
+	allModules := slices.Concat(artModules, platformModules, apexModules)
 	b.configuredModules = allModules
+	b.libraryToApex = maps.Clone(artModulesToApex)
+	maps.Copy(b.libraryToApex, platformModulesToApex)
+	maps.Copy(b.libraryToApex, apexModulesToApex)
 
 	// Do not add implLibModule to allModules as the impl lib is only used to collect the
 	// transitive source files
@@ -189,7 +198,7 @@
 	TransformResourcesToJar(ctx, srcjar, jarArgs, transitiveSrcFiles)
 
 	// Gather all the fragments dependencies.
-	b.fragments = gatherApexModulePairDepsWithTag(ctx, fragment)
+	b.fragments, b.apexNameToFragment = gatherFragments(ctx)
 
 	// Check the configuration of the boot modules.
 	// ART modules are checked by the art-bootclasspath-fragment.
@@ -198,7 +207,7 @@
 
 	b.generateClasspathProtoBuildActions(ctx)
 
-	bootDexJarByModule := b.generateHiddenAPIBuildActions(ctx, b.configuredModules, b.fragments)
+	bootDexJarByModule := b.generateHiddenAPIBuildActions(ctx, b.configuredModules, b.fragments, b.libraryToApex, b.apexNameToFragment)
 	buildRuleForBootJarsPackageCheck(ctx, bootDexJarByModule)
 
 	ctx.SetOutputFiles(android.Paths{b.hiddenAPIFlagsCSV}, "hiddenapi-flags.csv")
@@ -289,7 +298,8 @@
 }
 
 // generateHiddenAPIBuildActions generates all the hidden API related build rules.
-func (b *platformBootclasspathModule) generateHiddenAPIBuildActions(ctx android.ModuleContext, modules []android.Module, fragments []android.Module) bootDexJarByModule {
+func (b *platformBootclasspathModule) generateHiddenAPIBuildActions(ctx android.ModuleContext, modules []android.Module,
+	fragments []android.Module, libraryToApex map[android.Module]string, apexNameToFragment map[string]android.Module) bootDexJarByModule {
 	createEmptyHiddenApiFiles := func() {
 		paths := android.OutputPaths{b.hiddenAPIFlagsCSV, b.hiddenAPIIndexCSV, b.hiddenAPIMetadataCSV}
 		for _, path := range paths {
@@ -316,7 +326,7 @@
 	}
 
 	// Construct a list of ClasspathElement objects from the modules and fragments.
-	classpathElements := CreateClasspathElements(ctx, modules, fragments)
+	classpathElements := CreateClasspathElements(ctx, modules, fragments, libraryToApex, apexNameToFragment)
 
 	monolithicInfo := b.createAndProvideMonolithicHiddenAPIInfo(ctx, classpathElements)