2 files changed, 12 insertions, 20 deletions
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) {