diff --git a/apex/apex.go b/apex/apex.go
index 0481658..4d0e3f1 100644
--- a/apex/apex.go
+++ b/apex/apex.go
@@ -1127,6 +1127,9 @@
 	}
 	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 @@
 	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 e7d92c3..3daa4f8 100644
--- a/apex/prebuilt.go
+++ b/apex/prebuilt.go
@@ -301,13 +301,17 @@
 }
 
 // Implements android.DepInInSameApex
-func (p *prebuiltCommon) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
-	_, ok := tag.(exportedDependencyTag)
-	return ok
+func (m *prebuiltCommon) GetDepInSameApexChecker() android.DepInSameApexChecker {
+	return ApexPrebuiltDepInSameApexChecker{}
 }
 
-func (p *prebuiltCommon) IncomingDepIsInSameApex(tag blueprint.DependencyTag) bool {
-	return true
+type ApexPrebuiltDepInSameApexChecker struct {
+	android.BaseDepInSameApexChecker
+}
+
+func (m ApexPrebuiltDepInSameApexChecker) OutgoingDepIsInSameApex(tag blueprint.DependencyTag) bool {
+	_, ok := tag.(exportedDependencyTag)
+	return ok
 }
 
 func (p *prebuiltCommon) checkExportedDependenciesArePrebuilts(ctx android.ModuleContext) {