Merge "apexDepsMutator is a top-down mutator"
diff --git a/android/apex.go b/android/apex.go
index a4a5231..17ec9b1 100644
--- a/android/apex.go
+++ b/android/apex.go
@@ -38,9 +38,12 @@
 	Module
 	apexModuleBase() *ApexModuleBase
 
-	// Marks that this module should be built for the APEX of the specified name.
+	// Marks that this module should be built for the APEXes of the specified names.
 	// Call this before apex.apexMutator is run.
-	BuildForApex(apexName string)
+	BuildForApexes(apexNames []string)
+
+	// Returns the name of the APEXes that this modoule will be built for
+	ApexVariations() []string
 
 	// Returns the name of APEX that this module will be built for. Empty string
 	// is returned when 'IsForPlatform() == true'. Note that a module can be
@@ -66,7 +69,7 @@
 	IsInstallableToApex() bool
 
 	// Mutate this module into one or more variants each of which is built
-	// for an APEX marked via BuildForApex().
+	// for an APEX marked via BuildForApexes().
 	CreateApexVariations(mctx BottomUpMutatorContext) []Module
 
 	// Sets the name of the apex variant of this module. Called inside
@@ -110,14 +113,20 @@
 	return m
 }
 
-func (m *ApexModuleBase) BuildForApex(apexName string) {
+func (m *ApexModuleBase) BuildForApexes(apexNames []string) {
 	m.apexVariationsLock.Lock()
 	defer m.apexVariationsLock.Unlock()
-	if !InList(apexName, m.apexVariations) {
-		m.apexVariations = append(m.apexVariations, apexName)
+	for _, apexName := range apexNames {
+		if !InList(apexName, m.apexVariations) {
+			m.apexVariations = append(m.apexVariations, apexName)
+		}
 	}
 }
 
+func (m *ApexModuleBase) ApexVariations() []string {
+	return m.apexVariations
+}
+
 func (m *ApexModuleBase) ApexName() string {
 	return m.ApexProperties.ApexName
 }
@@ -218,18 +227,20 @@
 }
 
 // Update the map to mark that a module named moduleName is directly or indirectly
-// depended on by an APEX named apexName. Directly depending means that a module
+// depended on by the specified APEXes. Directly depending means that a module
 // is explicitly listed in the build definition of the APEX via properties like
 // native_shared_libs, java_libs, etc.
-func UpdateApexDependency(apexName string, moduleName string, directDep bool) {
+func UpdateApexDependency(apexNames []string, moduleName string, directDep bool) {
 	apexNamesMapMutex.Lock()
 	defer apexNamesMapMutex.Unlock()
-	apexNames, ok := apexNamesMap()[moduleName]
-	if !ok {
-		apexNames = make(map[string]bool)
-		apexNamesMap()[moduleName] = apexNames
+	for _, apexName := range apexNames {
+		apexesForModule, ok := apexNamesMap()[moduleName]
+		if !ok {
+			apexesForModule = make(map[string]bool)
+			apexNamesMap()[moduleName] = apexesForModule
+		}
+		apexesForModule[apexName] = apexesForModule[apexName] || directDep
 	}
-	apexNames[apexName] = apexNames[apexName] || directDep
 }
 
 // TODO(b/146393795): remove this when b/146393795 is fixed
diff --git a/apex/apex.go b/apex/apex.go
index e5f6424..1be87c3 100644
--- a/apex/apex.go
+++ b/apex/apex.go
@@ -1020,7 +1020,7 @@
 }
 
 func RegisterPostDepsMutators(ctx android.RegisterMutatorsContext) {
-	ctx.BottomUp("apex_deps", apexDepsMutator)
+	ctx.TopDown("apex_deps", apexDepsMutator)
 	ctx.BottomUp("apex", apexMutator).Parallel()
 	ctx.BottomUp("apex_flattened", apexFlattenedMutator).Parallel()
 	ctx.BottomUp("apex_uses", apexUsesMutator).Parallel()
@@ -1028,24 +1028,29 @@
 
 // Mark the direct and transitive dependencies of apex bundles so that they
 // can be built for the apex bundles.
-func apexDepsMutator(mctx android.BottomUpMutatorContext) {
+func apexDepsMutator(mctx android.TopDownMutatorContext) {
+	var apexBundleNames []string
+	var directDep bool
 	if a, ok := mctx.Module().(*apexBundle); ok && !a.vndkApex {
-		apexBundleName := mctx.ModuleName()
-		mctx.WalkDeps(func(child, parent android.Module) bool {
-			depName := mctx.OtherModuleName(child)
-			// If the parent is apexBundle, this child is directly depended.
-			_, directDep := parent.(*apexBundle)
-			android.UpdateApexDependency(apexBundleName, depName, directDep)
-
-			if am, ok := child.(android.ApexModule); ok && am.CanHaveApexVariants() &&
-				(directDep || am.DepIsInSameApex(mctx, child)) {
-				am.BuildForApex(apexBundleName)
-				return true
-			} else {
-				return false
-			}
-		})
+		apexBundleNames = []string{mctx.ModuleName()}
+		directDep = true
+	} else if am, ok := mctx.Module().(android.ApexModule); ok {
+		apexBundleNames = am.ApexVariations()
+		directDep = false
 	}
+
+	if len(apexBundleNames) == 0 {
+		return
+	}
+
+	mctx.VisitDirectDeps(func(child android.Module) {
+		depName := mctx.OtherModuleName(child)
+		if am, ok := child.(android.ApexModule); ok && am.CanHaveApexVariants() &&
+			(directDep || am.DepIsInSameApex(mctx, child)) {
+			android.UpdateApexDependency(apexBundleNames, depName, directDep)
+			am.BuildForApexes(apexBundleNames)
+		}
+	})
 }
 
 // Create apex variations if a module is included in APEX(s).