Support DepIsInSameApex using provider.

Bug: 377723687
Test: Unit tests and compare the ninja and mk files generated.
Change-Id: I8ec5e8a3a06d078602ebaa902cacb70987f1deda

12 files changed, 179 insertions, 74 deletions

diff --git a/android/apex.go b/android/apex.go
index 68d0ce8a2..780d091bf 100644
--- a/android/apex.go
+++ b/android/apex.go
@@ -155,11 +155,9 @@ type ApexBundleInfo struct {
 
 var ApexBundleInfoProvider = blueprint.NewMutatorProvider[ApexBundleInfo]("apex_mutate")
 
-// DepIsInSameApex defines an interface that should be used to determine whether a given dependency
-// should be considered as part of the same APEX as the current module or not. Note: this was
-// extracted from ApexModule to make it easier to define custom subsets of the ApexModule interface
-// and improve code navigation within the IDE.
-type DepIsInSameApex interface {
+// DepInSameApexChecker defines an interface that should be used to determine whether a given dependency
+// should be considered as part of the same APEX as the current module or not.
+type DepInSameApexChecker interface {
 	// OutgoingDepIsInSameApex tests if the module depended on via 'tag' is considered as part of
 	// the same APEX as this module. For example, a static lib dependency usually returns true here, while a
 	// shared lib dependency to a stub library returns false.
@@ -181,6 +179,15 @@ type DepIsInSameApex interface {
 	IncomingDepIsInSameApex(tag blueprint.DependencyTag) bool
 }
 
+// DepInSameApexInfo is a provider that wraps around a DepInSameApexChecker that can be
+// used to check if a dependency belongs to the same apex as the module when walking
+// through the dependencies of a module.
+type DepInSameApexInfo struct {
+	Checker DepInSameApexChecker
+}
+
+var DepInSameApexInfoProvider = blueprint.NewMutatorProvider[DepInSameApexInfo]("apex_unique")
+
 func IsDepInSameApex(ctx BaseModuleContext, module, dep Module) bool {
 	depTag := ctx.OtherModuleDependencyTag(dep)
 	if _, ok := depTag.(ExcludeFromApexContentsTag); ok {
@@ -189,12 +196,23 @@ func IsDepInSameApex(ctx BaseModuleContext, module, dep Module) bool {
 		return false
 	}
 
-	if m, ok := module.(DepIsInSameApex); ok && !m.OutgoingDepIsInSameApex(depTag) {
-		return false
+	if !ctx.EqualModules(ctx.Module(), module) {
+		if moduleInfo, ok := OtherModuleProvider(ctx, module, DepInSameApexInfoProvider); ok {
+			if !moduleInfo.Checker.OutgoingDepIsInSameApex(depTag) {
+				return false
+			}
+		}
+	} else {
+		if m, ok := ctx.Module().(ApexModule); ok && !m.GetDepInSameApexChecker().OutgoingDepIsInSameApex(depTag) {
+			return false
+		}
 	}
-	if d, ok := dep.(DepIsInSameApex); ok && !d.IncomingDepIsInSameApex(depTag) {
-		return false
+	if depInfo, ok := OtherModuleProvider(ctx, dep, DepInSameApexInfoProvider); ok {
+		if !depInfo.Checker.IncomingDepIsInSameApex(depTag) {
+			return false
+		}
 	}
+
 	return true
 }
 
@@ -213,7 +231,6 @@ func IsDepInSameApex(ctx BaseModuleContext, module, dep Module) bool {
 // mergedName) when the two APEXes have the same min_sdk_version requirement.
 type ApexModule interface {
 	Module
-	DepIsInSameApex
 
 	apexModuleBase() *ApexModuleBase
 
@@ -275,6 +292,8 @@ type ApexModule interface {
 	// deduping. This is turned on when, for example if use_apex_name_macro is set so that each
 	// apex variant should be built with different macro definitions.
 	UniqueApexVariations() bool
+
+	GetDepInSameApexChecker() DepInSameApexChecker
 }
 
 // Properties that are common to all module types implementing ApexModule interface.
@@ -331,7 +350,7 @@ func (m *ApexModuleBase) ApexTransitionMutatorSplit(ctx BaseModuleContext) []Ape
 }
 
 func (m *ApexModuleBase) ApexTransitionMutatorOutgoing(ctx OutgoingTransitionContext, info ApexInfo) ApexInfo {
-	if !ctx.Module().(DepIsInSameApex).OutgoingDepIsInSameApex(ctx.DepTag()) {
+	if !ctx.Module().(ApexModule).GetDepInSameApexChecker().OutgoingDepIsInSameApex(ctx.DepTag()) {
 		return ApexInfo{}
 	}
 	return info
@@ -343,7 +362,7 @@ func (m *ApexModuleBase) ApexTransitionMutatorIncoming(ctx IncomingTransitionCon
 		return ApexInfo{}
 	}
 
-	if !ctx.Module().(DepIsInSameApex).IncomingDepIsInSameApex(ctx.DepTag()) {
+	if !ctx.Module().(ApexModule).GetDepInSameApexChecker().IncomingDepIsInSameApex(ctx.DepTag()) {
 		return ApexInfo{}
 	}
 
@@ -460,19 +479,17 @@ func (m *ApexModuleBase) UniqueApexVariations() bool {
 }
 
 // Implements ApexModule
-func (m *ApexModuleBase) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
-	// By default, if there is a dependency from A to B, we try to include both in the same
-	// APEX, unless B is explicitly from outside of the APEX (i.e. a stubs lib). Thus, returning
-	// true. This is overridden by some module types like apex.ApexBundle, cc.Module,
-	// java.Module, etc.
+func (m *ApexModuleBase) GetDepInSameApexChecker() DepInSameApexChecker {
+	return BaseDepInSameApexChecker{}
+}
+
+type BaseDepInSameApexChecker struct{}
+
+func (m BaseDepInSameApexChecker) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
 	return true
 }
 
-func (m *ApexModuleBase) IncomingDepIsInSameApex(tag blueprint.DependencyTag) bool {
-	// By default, if there is a dependency from A to B, we try to include both in the same
-	// APEX, unless B is explicitly from outside of the APEX (i.e. a stubs lib). Thus, returning
-	// true. This is overridden by some module types like apex.ApexBundle, cc.Module,
-	// java.Module, etc.
+func (m BaseDepInSameApexChecker) IncomingDepIsInSameApex(tag blueprint.DependencyTag) bool {
 	return true
 }
 
diff --git a/apex/apex.go b/apex/apex.go
index 04816580d..4d0e3f183 100644
--- a/apex/apex.go
+++ b/apex/apex.go
@@ -1127,6 +1127,9 @@ func apexUniqueVariationsMutator(mctx android.BottomUpMutatorContext) {
 	}
 	if am, ok := mctx.Module().(android.ApexModule); ok {
 		android.UpdateUniqueApexVariationsForDeps(mctx, am)
+		android.SetProvider(mctx, android.DepInSameApexInfoProvider, android.DepInSameApexInfo{
+			Checker: am.GetDepInSameApexChecker(),
+		})
 	}
 }
 
@@ -1223,21 +1226,6 @@ const (
 	erofsFsType = "erofs"
 )
 
-var _ android.DepIsInSameApex = (*apexBundle)(nil)
-
-// Implements android.DepInInSameApex
-func (a *apexBundle) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
-	// direct deps of an APEX bundle are all part of the APEX bundle
-	// TODO(jiyong): shouldn't we look into the payload field of the dependencyTag?
-	return true
-}
-
-func (a *apexBundle) IncomingDepIsInSameApex(tag blueprint.DependencyTag) bool {
-	// direct deps of an APEX bundle are all part of the APEX bundle
-	// TODO(jiyong): shouldn't we look into the payload field of the dependencyTag?
-	return true
-}
-
 func (a *apexBundle) Exportable() bool {
 	return true
 }
diff --git a/apex/prebuilt.go b/apex/prebuilt.go
index e7d92c3a4..3daa4f81a 100644
--- a/apex/prebuilt.go
+++ b/apex/prebuilt.go
@@ -301,13 +301,17 @@ func (p *prebuiltCommon) prebuiltApexContentsDeps(ctx android.BottomUpMutatorCon
 }
 
 // Implements android.DepInInSameApex
-func (p *prebuiltCommon) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
-	_, ok := tag.(exportedDependencyTag)
-	return ok
+func (m *prebuiltCommon) GetDepInSameApexChecker() android.DepInSameApexChecker {
+	return ApexPrebuiltDepInSameApexChecker{}
+}
+
+type ApexPrebuiltDepInSameApexChecker struct {
+	android.BaseDepInSameApexChecker
 }
 
-func (p *prebuiltCommon) IncomingDepIsInSameApex(tag blueprint.DependencyTag) bool {
-	return true
+func (m ApexPrebuiltDepInSameApexChecker) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
+	_, ok := tag.(exportedDependencyTag)
+	return ok
 }
 
 func (p *prebuiltCommon) checkExportedDependenciesArePrebuilts(ctx android.ModuleContext) {
diff --git a/cc/cc.go b/cc/cc.go
index cb11fb845..36e336b94 100644
--- a/cc/cc.go
+++ b/cc/cc.go
@@ -4056,7 +4056,23 @@ func (c *Module) AndroidMkWriteAdditionalDependenciesForSourceAbiDiff(w io.Write
 var _ android.ApexModule = (*Module)(nil)
 
 // Implements android.ApexModule
-func (c *Module) OutgoingDepIsInSameApex(depTag blueprint.DependencyTag) bool {
+func (c *Module) GetDepInSameApexChecker() android.DepInSameApexChecker {
+	return CcDepInSameApexChecker{
+		Static:           c.static(),
+		HasStubsVariants: c.HasStubsVariants(),
+		IsLlndk:          c.IsLlndk(),
+		Host:             c.Host(),
+	}
+}
+
+type CcDepInSameApexChecker struct {
+	Static           bool
+	HasStubsVariants bool
+	IsLlndk          bool
+	Host             bool
+}
+
+func (c CcDepInSameApexChecker) OutgoingDepIsInSameApex(depTag blueprint.DependencyTag) bool {
 	if depTag == StubImplDepTag {
 		// We don't track from an implementation library to its stubs.
 		return false
@@ -4069,7 +4085,7 @@ func (c *Module) OutgoingDepIsInSameApex(depTag blueprint.DependencyTag) bool {
 	}
 
 	libDepTag, isLibDepTag := depTag.(libraryDependencyTag)
-	if isLibDepTag && c.static() && libDepTag.shared() {
+	if isLibDepTag && c.Static && libDepTag.shared() {
 		// shared_lib dependency from a static lib is considered as crossing
 		// the APEX boundary because the dependency doesn't actually is
 		// linked; the dependency is used only during the compilation phase.
@@ -4083,11 +4099,11 @@ func (c *Module) OutgoingDepIsInSameApex(depTag blueprint.DependencyTag) bool {
 	return true
 }
 
-func (c *Module) IncomingDepIsInSameApex(depTag blueprint.DependencyTag) bool {
-	if c.Host() {
+func (c CcDepInSameApexChecker) IncomingDepIsInSameApex(depTag blueprint.DependencyTag) bool {
+	if c.Host {
 		return false
 	}
-	if c.HasStubsVariants() {
+	if c.HasStubsVariants {
 		if IsSharedDepTag(depTag) && !IsExplicitImplSharedDepTag(depTag) {
 			// dynamic dep to a stubs lib crosses APEX boundary
 			return false
@@ -4100,7 +4116,7 @@ func (c *Module) IncomingDepIsInSameApex(depTag blueprint.DependencyTag) bool {
 			return false
 		}
 	}
-	if c.IsLlndk() {
+	if c.IsLlndk {
 		return false
 	}
 
diff --git a/java/aar.go b/java/aar.go
index 0a5a4c4d7..95387a358 100644
--- a/java/aar.go
+++ b/java/aar.go
@@ -1625,8 +1625,16 @@ var _ UsesLibraryDependency = (*AARImport)(nil)
 var _ android.ApexModule = (*AARImport)(nil)
 
 // Implements android.ApexModule
-func (a *AARImport) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
-	return a.depIsInSameApex(tag)
+func (m *AARImport) GetDepInSameApexChecker() android.DepInSameApexChecker {
+	return AARImportDepInSameApexChecker{}
+}
+
+type AARImportDepInSameApexChecker struct {
+	android.BaseDepInSameApexChecker
+}
+
+func (m AARImportDepInSameApexChecker) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
+	return depIsInSameApex(tag)
 }
 
 // Implements android.ApexModule
diff --git a/java/app.go b/java/app.go
index abbf034fa..ed2c35e74 100644
--- a/java/app.go
+++ b/java/app.go
@@ -1256,7 +1256,7 @@ func (a *AndroidApp) WalkPayloadDeps(ctx android.BaseModuleContext, do android.P
 		// TODO(ccross): Should this use android.DepIsInSameApex?  Right now it is applying the android app
 		// heuristics to every transitive dependency, when it should probably be using the heuristics of the
 		// immediate parent.
-		isExternal := !a.OutgoingDepIsInSameApex(ctx.OtherModuleDependencyTag(child))
+		isExternal := !a.GetDepInSameApexChecker().OutgoingDepIsInSameApex(ctx.OtherModuleDependencyTag(child))
 		if am, ok := child.(android.ApexModule); ok {
 			if !do(ctx, parent, am, isExternal) {
 				return false
@@ -1333,11 +1333,19 @@ func (a *AndroidApp) getCertString(ctx android.BaseModuleContext) string {
 	return a.overridableAppProperties.Certificate.GetOrDefault(ctx, "")
 }
 
-func (a *AndroidApp) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
+func (m *AndroidApp) GetDepInSameApexChecker() android.DepInSameApexChecker {
+	return AppDepInSameApexChecker{}
+}
+
+type AppDepInSameApexChecker struct {
+	android.BaseDepInSameApexChecker
+}
+
+func (m AppDepInSameApexChecker) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
 	if IsJniDepTag(tag) {
 		return true
 	}
-	return a.Library.OutgoingDepIsInSameApex(tag)
+	return depIsInSameApex(tag)
 }
 
 func (a *AndroidApp) Privileged() bool {
diff --git a/java/app_import.go b/java/app_import.go
index dc1aa740f..a122510de 100644
--- a/java/app_import.go
+++ b/java/app_import.go
@@ -628,7 +628,15 @@ func (a *AndroidAppImport) Privileged() bool {
 	return Bool(a.properties.Privileged)
 }
 
-func (a *AndroidAppImport) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
+func (m *AndroidAppImport) GetDepInSameApexChecker() android.DepInSameApexChecker {
+	return AppImportDepInSameApexChecker{}
+}
+
+type AppImportDepInSameApexChecker struct {
+	android.BaseDepInSameApexChecker
+}
+
+func (m AppImportDepInSameApexChecker) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
 	// android_app_import might have extra dependencies via uses_libs property.
 	// Don't track the dependency as we don't automatically add those libraries
 	// to the classpath. It should be explicitly added to java_libs property of APEX
diff --git a/java/base.go b/java/base.go
index 3a435a173..730554896 100644
--- a/java/base.go
+++ b/java/base.go
@@ -377,7 +377,7 @@ func (e *embeddableInModuleAndImport) initModuleAndImport(module android.Module)
 //
 // This cannot implement OutgoingDepIsInSameApex(...) directly as that leads to ambiguity with
 // the one provided by ApexModuleBase.
-func (e *embeddableInModuleAndImport) depIsInSameApex(tag blueprint.DependencyTag) bool {
+func depIsInSameApex(tag blueprint.DependencyTag) bool {
 	// dependencies other than the static linkage are all considered crossing APEX boundary
 	if tag == staticLibTag {
 		return true
@@ -2254,8 +2254,16 @@ func (j *Module) hasCode(ctx android.ModuleContext) bool {
 }
 
 // Implements android.ApexModule
-func (j *Module) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
-	return j.depIsInSameApex(tag)
+func (m *Module) GetDepInSameApexChecker() android.DepInSameApexChecker {
+	return JavaDepInSameApexChecker{}
+}
+
+type JavaDepInSameApexChecker struct {
+	android.BaseDepInSameApexChecker
+}
+
+func (m JavaDepInSameApexChecker) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
+	return depIsInSameApex(tag)
 }
 
 // Implements android.ApexModule
diff --git a/java/bootclasspath_fragment.go b/java/bootclasspath_fragment.go
index 65a05798a..8383a5a1e 100644
--- a/java/bootclasspath_fragment.go
+++ b/java/bootclasspath_fragment.go
@@ -410,7 +410,15 @@ func (i BootclasspathFragmentApexContentInfo) ProfileInstallPathInApex() string
 	return i.profileInstallPathInApex
 }
 
-func (b *BootclasspathFragmentModule) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
+func (m *BootclasspathFragmentModule) GetDepInSameApexChecker() android.DepInSameApexChecker {
+	return BootclasspathFragmentDepInSameApexChecker{}
+}
+
+type BootclasspathFragmentDepInSameApexChecker struct {
+	android.BaseDepInSameApexChecker
+}
+
+func (b BootclasspathFragmentDepInSameApexChecker) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
 	// If the module is a default module, do not check the tag
 	if tag == android.DefaultsDepTag {
 		return true
@@ -443,7 +451,7 @@ func (b *BootclasspathFragmentModule) OutgoingDepIsInSameApex(tag blueprint.Depe
 	if _, ok := tag.(hiddenAPIStubsDependencyTag); ok {
 		return false
 	}
-	panic(fmt.Errorf("boot_image module %q should not have a dependency tag %s", b, android.PrettyPrintTag(tag)))
+	panic(fmt.Errorf("boot_image module should not have a dependency tag %s", android.PrettyPrintTag(tag)))
 }
 
 func (b *BootclasspathFragmentModule) ShouldSupportSdkVersion(ctx android.BaseModuleContext, sdkVersion android.ApiLevel) error {
diff --git a/java/java.go b/java/java.go
index 91b7cf24b..82a9a62a6 100644
--- a/java/java.go
+++ b/java/java.go
@@ -3286,8 +3286,16 @@ func (j *Import) ClassLoaderContexts() dexpreopt.ClassLoaderContextMap {
 var _ android.ApexModule = (*Import)(nil)
 
 // Implements android.ApexModule
-func (j *Import) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
-	return j.depIsInSameApex(tag)
+func (m *Import) GetDepInSameApexChecker() android.DepInSameApexChecker {
+	return JavaImportDepInSameApexChecker{}
+}
+
+type JavaImportDepInSameApexChecker struct {
+	android.BaseDepInSameApexChecker
+}
+
+func (m JavaImportDepInSameApexChecker) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
+	return depIsInSameApex(tag)
 }
 
 // Implements android.ApexModule
diff --git a/java/sdk_library.go b/java/sdk_library.go
index b6bac2d96..fafb44801 100644
--- a/java/sdk_library.go
+++ b/java/sdk_library.go
@@ -1708,14 +1708,22 @@ func (module *SdkLibrary) compareAgainstLatestApi(apiScope *apiScope) bool {
 }
 
 // Implements android.ApexModule
-func (module *SdkLibrary) OutgoingDepIsInSameApex(depTag blueprint.DependencyTag) bool {
-	if depTag == xmlPermissionsFileTag {
+func (m *SdkLibrary) GetDepInSameApexChecker() android.DepInSameApexChecker {
+	return SdkLibraryDepInSameApexChecker{}
+}
+
+type SdkLibraryDepInSameApexChecker struct {
+	android.BaseDepInSameApexChecker
+}
+
+func (m SdkLibraryDepInSameApexChecker) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
+	if tag == xmlPermissionsFileTag {
 		return true
 	}
-	if depTag == implLibraryTag {
+	if tag == implLibraryTag {
 		return true
 	}
-	return module.Library.OutgoingDepIsInSameApex(depTag)
+	return depIsInSameApex(tag)
 }
 
 // Implements android.ApexModule
@@ -2126,8 +2134,16 @@ func (module *SdkLibraryImport) DepsMutator(ctx android.BottomUpMutatorContext)
 var _ android.ApexModule = (*SdkLibraryImport)(nil)
 
 // Implements android.ApexModule
-func (module *SdkLibraryImport) OutgoingDepIsInSameApex(depTag blueprint.DependencyTag) bool {
-	if depTag == xmlPermissionsFileTag {
+func (m *SdkLibraryImport) GetDepInSameApexChecker() android.DepInSameApexChecker {
+	return SdkLibraryImportDepIsInSameApexChecker{}
+}
+
+type SdkLibraryImportDepIsInSameApexChecker struct {
+	android.BaseDepInSameApexChecker
+}
+
+func (m SdkLibraryImportDepIsInSameApexChecker) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
+	if tag == xmlPermissionsFileTag {
 		return true
 	}
 
diff --git a/rust/rust.go b/rust/rust.go
index 4eec5d254..4fd800282 100644
--- a/rust/rust.go
+++ b/rust/rust.go
@@ -2115,12 +2115,28 @@ func (mod *Module) AlwaysRequiresPlatformApexVariant() bool {
 }
 
 // Implements android.ApexModule
-func (mod *Module) OutgoingDepIsInSameApex(depTag blueprint.DependencyTag) bool {
+type RustDepInSameApexChecker struct {
+	Static           bool
+	HasStubsVariants bool
+	ApexExclude      bool
+	Host             bool
+}
+
+func (mod *Module) GetDepInSameApexChecker() android.DepInSameApexChecker {
+	return RustDepInSameApexChecker{
+		Static:           mod.Static(),
+		HasStubsVariants: mod.HasStubsVariants(),
+		ApexExclude:      mod.ApexExclude(),
+		Host:             mod.Host(),
+	}
+}
+
+func (r RustDepInSameApexChecker) OutgoingDepIsInSameApex(depTag blueprint.DependencyTag) bool {
 	if depTag == procMacroDepTag || depTag == customBindgenDepTag {
 		return false
 	}
 
-	if mod.Static() && cc.IsSharedDepTag(depTag) {
+	if r.Static && cc.IsSharedDepTag(depTag) {
 		// shared_lib dependency from a static lib is considered as crossing
 		// the APEX boundary because the dependency doesn't actually is
 		// linked; the dependency is used only during the compilation phase.
@@ -2137,23 +2153,23 @@ func (mod *Module) OutgoingDepIsInSameApex(depTag blueprint.DependencyTag) bool
 	}
 
 	// TODO(b/362509506): remove once all apex_exclude uses are switched to stubs.
-	if mod.ApexExclude() {
+	if r.ApexExclude {
 		return false
 	}
 
 	return true
 }
 
-func (mod *Module) IncomingDepIsInSameApex(depTag blueprint.DependencyTag) bool {
-	if mod.Host() {
+func (r RustDepInSameApexChecker) IncomingDepIsInSameApex(depTag blueprint.DependencyTag) bool {
+	if r.Host {
 		return false
 	}
 	// TODO(b/362509506): remove once all apex_exclude uses are switched to stubs.
-	if mod.ApexExclude() {
+	if r.ApexExclude {
 		return false
 	}
 
-	if mod.HasStubsVariants() {
+	if r.HasStubsVariants {
 		if cc.IsSharedDepTag(depTag) && !cc.IsExplicitImplSharedDepTag(depTag) {
 			// dynamic dep to a stubs lib crosses APEX boundary
 			return false