Convert dex2oatPathFromDep to use ModuleProxy.

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

6 files changed, 23 insertions, 16 deletions

diff --git a/android/module.go b/android/module.go
index c88d04556..b9489b482 100644
--- a/android/module.go
+++ b/android/module.go
@@ -1891,15 +1891,16 @@ var CommonModuleInfoKey = blueprint.NewProvider[CommonModuleInfo]()
 
 type PrebuiltModuleInfo struct {
 	SourceExists bool
+	UsePrebuilt  bool
 }
 
 var PrebuiltModuleInfoProvider = blueprint.NewProvider[PrebuiltModuleInfo]()
 
-type HostToolProviderData struct {
+type HostToolProviderInfo struct {
 	HostToolPath OptionalPath
 }
 
-var HostToolProviderKey = blueprint.NewProvider[HostToolProviderData]()
+var HostToolProviderInfoProvider = blueprint.NewProvider[HostToolProviderInfo]()
 
 type SourceFileGenerator interface {
 	GeneratedSourceFiles() Paths
@@ -2212,10 +2213,11 @@ func (m *ModuleBase) GenerateBuildActions(blueprintCtx blueprint.ModuleContext)
 	if p, ok := m.module.(PrebuiltInterface); ok && p.Prebuilt() != nil {
 		SetProvider(ctx, PrebuiltModuleInfoProvider, PrebuiltModuleInfo{
 			SourceExists: p.Prebuilt().SourceExists(),
+			UsePrebuilt:  p.Prebuilt().UsePrebuilt(),
 		})
 	}
 	if h, ok := m.module.(HostToolProvider); ok {
-		SetProvider(ctx, HostToolProviderKey, HostToolProviderData{
+		SetProvider(ctx, HostToolProviderInfoProvider, HostToolProviderInfo{
 			HostToolPath: h.HostToolPath()})
 	}
 
diff --git a/android/module_proxy.go b/android/module_proxy.go
index 8cc8fa1e7..afca0d72b 100644
--- a/android/module_proxy.go
+++ b/android/module_proxy.go
@@ -11,6 +11,10 @@ type ModuleProxy struct {
 
 var _ Module = (*ModuleProxy)(nil)
 
+func (m ModuleProxy) IsNil() bool {
+	return m.module.IsNil()
+}
+
 func (m ModuleProxy) Name() string {
 	return m.module.Name()
 }
diff --git a/android/proto.go b/android/proto.go
index 66faa20ac..91d67322b 100644
--- a/android/proto.go
+++ b/android/proto.go
@@ -75,7 +75,7 @@ func GetProtoFlags(ctx ModuleContext, p *ProtoProperties) ProtoFlags {
 	}
 
 	ctx.VisitDirectDepsProxyWithTag(ProtoPluginDepTag, func(dep ModuleProxy) {
-		if h, ok := OtherModuleProvider(ctx, dep, HostToolProviderKey); !ok || !h.HostToolPath.Valid() {
+		if h, ok := OtherModuleProvider(ctx, dep, HostToolProviderInfoProvider); !ok || !h.HostToolPath.Valid() {
 			ctx.PropertyErrorf("proto.plugin", "module %q is not a host tool provider",
 				ctx.OtherModuleName(dep))
 		} else {
diff --git a/cc/cc.go b/cc/cc.go
index 16aa4a792..b525ccb81 100644
--- a/cc/cc.go
+++ b/cc/cc.go
@@ -2384,7 +2384,7 @@ func CreateCommonLinkableInfo(ctx android.ModuleContext, mod VersionedLinkableIn
 		RelativeInstallPath:  mod.RelativeInstallPath(),
 		// TODO(b/362509506): remove this once all apex_exclude uses are switched to stubs.
 		RustApexExclude: mod.RustApexExclude(),
-		Bootstrap:            mod.Bootstrap(),
+		Bootstrap:       mod.Bootstrap(),
 	}
 }
 
diff --git a/dexpreopt/config.go b/dexpreopt/config.go
index ee23b5121..c5cafb1cf 100644
--- a/dexpreopt/config.go
+++ b/dexpreopt/config.go
@@ -507,36 +507,37 @@ func dex2oatPathFromDep(ctx android.ModuleContext) android.Path {
 	// final variants, while prebuilt_postdeps needs to run before many of the
 	// PostDeps mutators, like the APEX mutators). Hence we need to dig out the
 	// prebuilt explicitly here instead.
-	var dex2oatModule android.Module
-	ctx.WalkDeps(func(child, parent android.Module) bool {
-		if parent == ctx.Module() && ctx.OtherModuleDependencyTag(child) == Dex2oatDepTag {
+	var dex2oatModule android.ModuleProxy
+	ctx.WalkDepsProxy(func(child, parent android.ModuleProxy) bool {
+		prebuiltInfo, isPrebuilt := android.OtherModuleProvider(ctx, child, android.PrebuiltModuleInfoProvider)
+		if ctx.EqualModules(parent, ctx.Module()) && ctx.OtherModuleDependencyTag(child) == Dex2oatDepTag {
 			// Found the source module, or prebuilt module that has replaced the source.
 			dex2oatModule = child
-			if android.IsModulePrebuilt(child) {
+			if isPrebuilt {
 				return false // If it's the prebuilt we're done.
 			} else {
 				return true // Recurse to check if the source has a prebuilt dependency.
 			}
 		}
-		if parent == dex2oatModule && ctx.OtherModuleDependencyTag(child) == android.PrebuiltDepTag {
-			if p := android.GetEmbeddedPrebuilt(child); p != nil && p.UsePrebuilt() {
+		if ctx.EqualModules(parent, dex2oatModule) && ctx.OtherModuleDependencyTag(child) == android.PrebuiltDepTag {
+			if isPrebuilt && prebuiltInfo.UsePrebuilt {
 				dex2oatModule = child // Found a prebuilt that should be used.
 			}
 		}
 		return false
 	})
 
-	if dex2oatModule == nil {
+	if dex2oatModule.IsNil() {
 		// If this happens there's probably a missing call to AddToolDeps in DepsMutator.
 		panic(fmt.Sprintf("Failed to lookup %s dependency", dex2oatBin))
 	}
 
-	dex2oatPath := dex2oatModule.(android.HostToolProvider).HostToolPath()
-	if !dex2oatPath.Valid() {
+	dex2oatPath, ok := android.OtherModuleProvider(ctx, dex2oatModule, android.HostToolProviderInfoProvider)
+	if !ok || !dex2oatPath.HostToolPath.Valid() {
 		panic(fmt.Sprintf("Failed to find host tool path in %s", dex2oatModule))
 	}
 
-	return dex2oatPath.Path()
+	return dex2oatPath.HostToolPath.Path()
 }
 
 // createGlobalSoongConfig creates a GlobalSoongConfig from the current context.
diff --git a/genrule/genrule.go b/genrule/genrule.go
index 9a8524bc7..65f74ce0f 100644
--- a/genrule/genrule.go
+++ b/genrule/genrule.go
@@ -349,7 +349,7 @@ func (g *Module) generateCommonBuildActions(ctx android.ModuleContext) {
 				// replaced the dependency.
 				module := android.PrebuiltGetPreferred(ctx, proxy)
 				tool := ctx.OtherModuleName(module)
-				if h, ok := android.OtherModuleProvider(ctx, module, android.HostToolProviderKey); ok {
+				if h, ok := android.OtherModuleProvider(ctx, module, android.HostToolProviderInfoProvider); ok {
 					// A HostToolProvider provides the path to a tool, which will be copied
 					// into the sandbox.
 					if !android.OtherModuleProviderOrDefault(ctx, module, android.CommonModuleInfoKey).Enabled {