Convert apex mutator to a TransitionMutator

Replace the apex mutator with a TransitionMutator.  Requires moving the
base.apexInfo value into a provider so that it is still present for the
Mutate pass.

Test: go test ./...
Test: no change to out/soong/build-${TARGET_PRODUCT}.ninja
Change-Id: I1c898eaf30b4021f0f96f439cc0b3b3173710fc7
diff --git a/android/apex.go b/android/apex.go
index c0acada..dc0aeed 100644
--- a/android/apex.go
+++ b/android/apex.go
@@ -90,7 +90,13 @@
 	TestApexes []string
 }
 
-var ApexInfoProvider = blueprint.NewMutatorProvider[ApexInfo]("apex")
+// AllApexInfo holds the ApexInfo of all apexes that include this module.
+type AllApexInfo struct {
+	ApexInfos []ApexInfo
+}
+
+var ApexInfoProvider = blueprint.NewMutatorProvider[ApexInfo]("apex_mutate")
+var AllApexInfoProvider = blueprint.NewMutatorProvider[*AllApexInfo]("apex_info")
 
 func (i ApexInfo) AddJSONData(d *map[string]interface{}) {
 	(*d)["Apex"] = map[string]interface{}{
@@ -586,75 +592,131 @@
 	return merged, aliases
 }
 
-// CreateApexVariations mutates a given module into multiple apex variants each of which is for an
-// apexBundle (and/or the platform) where the module is part of.
-func CreateApexVariations(mctx BottomUpMutatorContext, module ApexModule) []Module {
+// IncomingApexTransition is called by apexTransitionMutator.IncomingTransition on modules that can be in apexes.
+// The incomingVariation can be either the name of an apex if the dependency is coming directly from an apex
+// module, or it can be the name of an apex variation (e.g. apex10000) if it is coming from another module that
+// is in the apex.
+func IncomingApexTransition(ctx IncomingTransitionContext, incomingVariation string) string {
+	module := ctx.Module().(ApexModule)
 	base := module.apexModuleBase()
 
+	var apexInfos []ApexInfo
+	if allApexInfos, ok := ModuleProvider(ctx, AllApexInfoProvider); ok {
+		apexInfos = allApexInfos.ApexInfos
+	}
+
+	// Dependencies from platform variations go to the platform variation.
+	if incomingVariation == "" {
+		return ""
+	}
+
+	// If this module has no apex variations the use the platform variation.
+	if len(apexInfos) == 0 {
+		return ""
+	}
+
+	// Convert the list of apex infos into from the AllApexInfoProvider into the merged list
+	// of apex variations and the aliases from apex names to apex variations.
+	var aliases [][2]string
+	if !module.UniqueApexVariations() && !base.ApexProperties.UniqueApexVariationsForDeps {
+		apexInfos, aliases = mergeApexVariations(apexInfos)
+	}
+
+	// Check if the incoming variation matches an apex name, and if so use the corresponding
+	// apex variation.
+	aliasIndex := slices.IndexFunc(aliases, func(alias [2]string) bool {
+		return alias[0] == incomingVariation
+	})
+	if aliasIndex >= 0 {
+		return aliases[aliasIndex][1]
+	}
+
+	// Check if the incoming variation matches an apex variation.
+	apexIndex := slices.IndexFunc(apexInfos, func(info ApexInfo) bool {
+		return info.ApexVariationName == incomingVariation
+	})
+	if apexIndex >= 0 {
+		return incomingVariation
+	}
+
+	return ""
+}
+
+func MutateApexTransition(ctx BaseModuleContext, variation string) {
+	module := ctx.Module().(ApexModule)
+	base := module.apexModuleBase()
+	platformVariation := variation == ""
+
+	var apexInfos []ApexInfo
+	if allApexInfos, ok := ModuleProvider(ctx, AllApexInfoProvider); ok {
+		apexInfos = allApexInfos.ApexInfos
+	}
+
 	// Shortcut
-	if len(base.apexInfos) == 0 {
-		return nil
+	if len(apexInfos) == 0 {
+		return
 	}
 
 	// Do some validity checks.
 	// TODO(jiyong): is this the right place?
-	base.checkApexAvailableProperty(mctx)
+	base.checkApexAvailableProperty(ctx)
 
-	apexInfos := base.apexInfos
-	// base.apexInfos is only needed to propagate the list of apexes from apexInfoMutator to
-	// apexMutator. It is no longer accurate after mergeApexVariations, and won't be copied to
-	// all but the first created variant. Clear it so it doesn't accidentally get used later.
-	base.apexInfos = nil
-
-	slices.SortFunc(apexInfos, func(a, b ApexInfo) int {
-		return strings.Compare(a.ApexVariationName, b.ApexVariationName)
-	})
-
-	var aliases [][2]string
-	if !mctx.Module().(ApexModule).UniqueApexVariations() && !base.ApexProperties.UniqueApexVariationsForDeps {
-		apexInfos, aliases = mergeApexVariations(apexInfos)
+	if !module.UniqueApexVariations() && !base.ApexProperties.UniqueApexVariationsForDeps {
+		apexInfos, _ = mergeApexVariations(apexInfos)
 	}
 
 	var inApex ApexMembership
 	for _, a := range apexInfos {
 		for _, apexContents := range a.ApexContents {
-			inApex = inApex.merge(apexContents.contents[mctx.ModuleName()])
+			inApex = inApex.merge(apexContents.contents[ctx.ModuleName()])
 		}
 	}
 	base.ApexProperties.InAnyApex = true
 	base.ApexProperties.DirectlyInAnyApex = inApex == directlyInApex
 
-	defaultVariation := ""
-	mctx.SetDefaultDependencyVariation(&defaultVariation)
+	if platformVariation && !ctx.Host() && !module.AvailableFor(AvailableToPlatform) {
+		// Do not install the module for platform, but still allow it to output
+		// uninstallable AndroidMk entries in certain cases when they have side
+		// effects.  TODO(jiyong): move this routine to somewhere else
+		module.MakeUninstallable()
+	}
+	if !platformVariation {
+		var thisApexInfo ApexInfo
 
-	variations := []string{defaultVariation}
-	testApexes := []string{}
+		apexIndex := slices.IndexFunc(apexInfos, func(info ApexInfo) bool {
+			return info.ApexVariationName == variation
+		})
+		if apexIndex >= 0 {
+			thisApexInfo = apexInfos[apexIndex]
+		} else {
+			panic(fmt.Errorf("failed to find apexInfo for incoming variation %q", variation))
+		}
+
+		SetProvider(ctx, ApexInfoProvider, thisApexInfo)
+	}
+
+	// Set the value of TestApexes in every single apex variant.
+	// This allows each apex variant to be aware of the test apexes in the user provided apex_available.
+	var testApexes []string
 	for _, a := range apexInfos {
-		variations = append(variations, a.ApexVariationName)
 		testApexes = append(testApexes, a.TestApexes...)
 	}
-	modules := mctx.CreateVariations(variations...)
-	for i, mod := range modules {
-		platformVariation := i == 0
-		if platformVariation && !mctx.Host() && !mod.(ApexModule).AvailableFor(AvailableToPlatform) {
-			// Do not install the module for platform, but still allow it to output
-			// uninstallable AndroidMk entries in certain cases when they have side
-			// effects.  TODO(jiyong): move this routine to somewhere else
-			mod.MakeUninstallable()
-		}
-		if !platformVariation {
-			mctx.SetVariationProvider(mod, ApexInfoProvider, apexInfos[i-1])
-		}
-		// Set the value of TestApexes in every single apex variant.
-		// This allows each apex variant to be aware of the test apexes in the user provided apex_available.
-		mod.(ApexModule).apexModuleBase().ApexProperties.TestApexes = testApexes
-	}
+	base.ApexProperties.TestApexes = testApexes
 
-	for _, alias := range aliases {
-		mctx.CreateAliasVariation(alias[0], alias[1])
-	}
+}
 
-	return modules
+func ApexInfoMutator(ctx TopDownMutatorContext, module ApexModule) {
+	base := module.apexModuleBase()
+	if len(base.apexInfos) > 0 {
+		apexInfos := slices.Clone(base.apexInfos)
+		slices.SortFunc(apexInfos, func(a, b ApexInfo) int {
+			return strings.Compare(a.ApexVariationName, b.ApexVariationName)
+		})
+		SetProvider(ctx, AllApexInfoProvider, &AllApexInfo{apexInfos})
+		// base.apexInfos is only needed to propagate the list of apexes from the apex module to its
+		// contents within apexInfoMutator. Clear it so it doesn't accidentally get used later.
+		base.apexInfos = nil
+	}
 }
 
 // UpdateUniqueApexVariationsForDeps sets UniqueApexVariationsForDeps if any dependencies that are
@@ -665,13 +727,16 @@
 	// InApexVariants list in common. It is used instead of DepIsInSameApex because it needs to
 	// determine if the dep is in the same APEX due to being directly included, not only if it
 	// is included _because_ it is a dependency.
-	anyInSameApex := func(a, b []ApexInfo) bool {
-		collectApexes := func(infos []ApexInfo) []string {
-			var ret []string
-			for _, info := range infos {
-				ret = append(ret, info.InApexVariants...)
+	anyInSameApex := func(a, b ApexModule) bool {
+		collectApexes := func(m ApexModule) []string {
+			if allApexInfo, ok := OtherModuleProvider(mctx, m, AllApexInfoProvider); ok {
+				var ret []string
+				for _, info := range allApexInfo.ApexInfos {
+					ret = append(ret, info.InApexVariants...)
+				}
+				return ret
 			}
-			return ret
+			return nil
 		}
 
 		aApexes := collectApexes(a)
@@ -689,7 +754,7 @@
 	// If any of the dependencies requires unique apex variations, so does this module.
 	mctx.VisitDirectDeps(func(dep Module) {
 		if depApexModule, ok := dep.(ApexModule); ok {
-			if anyInSameApex(depApexModule.apexModuleBase().apexInfos, am.apexModuleBase().apexInfos) &&
+			if anyInSameApex(depApexModule, am) &&
 				(depApexModule.UniqueApexVariations() ||
 					depApexModule.apexModuleBase().ApexProperties.UniqueApexVariationsForDeps) {
 				am.apexModuleBase().ApexProperties.UniqueApexVariationsForDeps = true
diff --git a/apex/apex.go b/apex/apex.go
index cb8449c..ef57d7e 100644
--- a/apex/apex.go
+++ b/apex/apex.go
@@ -73,7 +73,7 @@
 	// Run mark_platform_availability before the apexMutator as the apexMutator needs to know whether
 	// it should create a platform variant.
 	ctx.BottomUp("mark_platform_availability", markPlatformAvailability).Parallel()
-	ctx.BottomUp("apex", apexMutator).Parallel()
+	ctx.Transition("apex", &apexTransitionMutator{})
 	ctx.BottomUp("apex_directly_in_any", apexDirectlyInAnyMutator).Parallel()
 	ctx.BottomUp("apex_dcla_deps", apexDCLADepsMutator).Parallel()
 	// Register after apex_info mutator so that it can use ApexVariationName
@@ -1084,6 +1084,10 @@
 	if a, ok := mctx.Module().(ApexInfoMutator); ok {
 		a.ApexInfoMutator(mctx)
 	}
+
+	if am, ok := mctx.Module().(android.ApexModule); ok {
+		android.ApexInfoMutator(mctx, am)
+	}
 	enforceAppUpdatability(mctx)
 }
 
@@ -1284,40 +1288,41 @@
 	}
 }
 
-// apexMutator visits each module and creates apex variations if the module was marked in the
-// previous run of apexInfoMutator.
-func apexMutator(mctx android.BottomUpMutatorContext) {
-	if !mctx.Module().Enabled() {
-		return
-	}
+type apexTransitionMutator struct{}
 
-	// This is the usual path.
-	if am, ok := mctx.Module().(android.ApexModule); ok && am.CanHaveApexVariants() {
-		android.CreateApexVariations(mctx, am)
-		return
-	}
-
+func (a *apexTransitionMutator) Split(ctx android.BaseModuleContext) []string {
 	// apexBundle itself is mutated so that it and its dependencies have the same apex variant.
-	// Note that a default variation "" is also created as an alias, and the default dependency
-	// variation is set to the default variation. This is to allow an apex to depend on another
-	// module which is outside of the apex. This is because the dependent module is not mutated
-	// for this apex variant.
-	createApexVariation := func(apexBundleName string) {
-		defaultVariation := ""
-		mctx.SetDefaultDependencyVariation(&defaultVariation)
-		mctx.CreateVariations(apexBundleName)
-		mctx.CreateAliasVariation(defaultVariation, apexBundleName)
-	}
-
-	if ai, ok := mctx.Module().(ApexInfoMutator); ok && apexModuleTypeRequiresVariant(ai) {
-		createApexVariation(ai.ApexVariationName())
-	} else if o, ok := mctx.Module().(*OverrideApex); ok {
+	if ai, ok := ctx.Module().(ApexInfoMutator); ok && apexModuleTypeRequiresVariant(ai) {
+		return []string{ai.ApexVariationName()}
+	} else if o, ok := ctx.Module().(*OverrideApex); ok {
 		apexBundleName := o.GetOverriddenModuleName()
 		if apexBundleName == "" {
-			mctx.ModuleErrorf("base property is not set")
-			return
+			ctx.ModuleErrorf("base property is not set")
 		}
-		createApexVariation(apexBundleName)
+		return []string{apexBundleName}
+	}
+	return []string{""}
+}
+
+func (a *apexTransitionMutator) OutgoingTransition(ctx android.OutgoingTransitionContext, sourceVariation string) string {
+	return sourceVariation
+}
+
+func (a *apexTransitionMutator) IncomingTransition(ctx android.IncomingTransitionContext, incomingVariation string) string {
+	if am, ok := ctx.Module().(android.ApexModule); ok && am.CanHaveApexVariants() {
+		return android.IncomingApexTransition(ctx, incomingVariation)
+	} else if ai, ok := ctx.Module().(ApexInfoMutator); ok {
+		return ai.ApexVariationName()
+	} else if o, ok := ctx.Module().(*OverrideApex); ok {
+		return o.GetOverriddenModuleName()
+	}
+
+	return ""
+}
+
+func (a *apexTransitionMutator) Mutate(ctx android.BottomUpMutatorContext, variation string) {
+	if am, ok := ctx.Module().(android.ApexModule); ok && am.CanHaveApexVariants() {
+		android.MutateApexTransition(ctx, variation)
 	}
 }
 
diff --git a/java/testing.go b/java/testing.go
index 631d516..5ae326d 100644
--- a/java/testing.go
+++ b/java/testing.go
@@ -711,7 +711,7 @@
 
 func registerFakeApexMutator(ctx android.RegistrationContext) {
 	ctx.PostDepsMutators(func(ctx android.RegisterMutatorsContext) {
-		ctx.BottomUp("apex", fakeApexMutator).Parallel()
+		ctx.Transition("apex", &fakeApexMutator{})
 	})
 }
 
@@ -726,16 +726,30 @@
 // `apex_available`. It helps us avoid a dependency on the real mutator defined in "soong-apex",
 // which will cause a cyclic dependency, and it provides an easy way to create an APEX variant for
 // testing without dealing with all the complexities in the real mutator.
-func fakeApexMutator(mctx android.BottomUpMutatorContext) {
-	switch mctx.Module().(type) {
+type fakeApexMutator struct{}
+
+func (f *fakeApexMutator) Split(ctx android.BaseModuleContext) []string {
+	switch ctx.Module().(type) {
 	case *Library, *SdkLibrary:
-		if len(mctx.Module().(apexModuleBase).ApexAvailable()) > 0 {
-			modules := mctx.CreateVariations("", "apex1000")
-			apexInfo := android.ApexInfo{
-				ApexVariationName: "apex1000",
-			}
-			mctx.SetVariationProvider(modules[1], android.ApexInfoProvider, apexInfo)
+		return []string{"", "apex1000"}
+	}
+	return []string{""}
+}
+
+func (f *fakeApexMutator) OutgoingTransition(ctx android.OutgoingTransitionContext, sourceVariation string) string {
+	return sourceVariation
+}
+
+func (f *fakeApexMutator) IncomingTransition(ctx android.IncomingTransitionContext, incomingVariation string) string {
+	return incomingVariation
+}
+
+func (f *fakeApexMutator) Mutate(ctx android.BottomUpMutatorContext, variation string) {
+	if variation != "" {
+		apexInfo := android.ApexInfo{
+			ApexVariationName: "apex1000",
 		}
+		android.SetProvider(ctx, android.ApexInfoProvider, apexInfo)
 	}
 }