Add a few module visiting methods that return ModuleProxy.

Change a few places to use these new methods when visiting dependencies.

Bug: 358427516
Test: Manual verified the generated mk and ninja files.
Change-Id: I7426843d76d79d41eb60fce5e796c14f968a0f1d
diff --git a/android/base_module_context.go b/android/base_module_context.go
index 670537f..e24ce9d 100644
--- a/android/base_module_context.go
+++ b/android/base_module_context.go
@@ -136,12 +136,14 @@
 	// multiple direct dependencies on the same module visit will be called multiple times on
 	// that module and OtherModuleDependencyTag will return a different tag for each.
 	//
-	// The Module passed to the visit function should not be retained outside of the visit function, it may be
+	// The ModuleProxy passed to the visit function should not be retained outside of the visit function, it may be
 	// invalidated by future mutators.
-	VisitDirectDepsProxyAllowDisabled(visit func(proxy Module))
+	VisitDirectDepsProxyAllowDisabled(visit func(proxy ModuleProxy))
 
 	VisitDirectDepsWithTag(tag blueprint.DependencyTag, visit func(Module))
 
+	VisitDirectDepsProxyWithTag(tag blueprint.DependencyTag, visit func(proxy ModuleProxy))
+
 	// VisitDirectDepsIf calls pred for each direct dependency, and if pred returns true calls visit.  If there are
 	// multiple direct dependencies on the same module pred and visit will be called multiple times on that module and
 	// OtherModuleDependencyTag will return a different tag for each.  It skips any
@@ -173,7 +175,7 @@
 	//
 	// The Modules passed to the visit function should not be retained outside of the visit function, they may be
 	// invalidated by future mutators.
-	WalkDepsProxy(visit func(child, parent Module) bool)
+	WalkDepsProxy(visit func(child, parent ModuleProxy) bool)
 
 	// GetWalkPath is supposed to be called in visit function passed in WalkDeps()
 	// and returns a top-down dependency path from a start module to current child module.
@@ -314,6 +316,7 @@
 type AllowDisabledModuleDependency interface {
 	blueprint.DependencyTag
 	AllowDisabledModuleDependency(target Module) bool
+	AllowDisabledModuleDependencyProxy(ctx OtherModuleProviderContext, target ModuleProxy) bool
 }
 
 type AlwaysAllowDisabledModuleDependencyTag struct{}
@@ -322,6 +325,10 @@
 	return true
 }
 
+func (t AlwaysAllowDisabledModuleDependencyTag) AllowDisabledModuleDependencyProxy(OtherModuleProviderContext, ModuleProxy) bool {
+	return true
+}
+
 func (b *baseModuleContext) validateAndroidModule(module blueprint.Module, tag blueprint.DependencyTag, strict bool) Module {
 	aModule, _ := module.(Module)
 
@@ -346,6 +353,28 @@
 	return aModule
 }
 
+func (b *baseModuleContext) validateAndroidModuleProxy(
+	module blueprint.ModuleProxy, tag blueprint.DependencyTag, strict bool) *ModuleProxy {
+	aModule := ModuleProxy{module: module}
+
+	if !strict {
+		return &aModule
+	}
+
+	if !OtherModuleProviderOrDefault(b, module, CommonPropertiesProviderKey).Enabled {
+		if t, ok := tag.(AllowDisabledModuleDependency); !ok || !t.AllowDisabledModuleDependencyProxy(b, aModule) {
+			if b.Config().AllowMissingDependencies() {
+				b.AddMissingDependencies([]string{b.OtherModuleName(aModule)})
+			} else {
+				b.ModuleErrorf("depends on disabled module %q", b.OtherModuleName(aModule))
+			}
+		}
+		return nil
+	}
+
+	return &aModule
+}
+
 type dep struct {
 	mod blueprint.Module
 	tag blueprint.DependencyTag
@@ -426,7 +455,7 @@
 	})
 }
 
-func (b *baseModuleContext) VisitDirectDepsProxyAllowDisabled(visit func(proxy Module)) {
+func (b *baseModuleContext) VisitDirectDepsProxyAllowDisabled(visit func(proxy ModuleProxy)) {
 	b.bp.VisitDirectDepsProxy(func(module blueprint.ModuleProxy) {
 		visit(ModuleProxy{
 			module: module,
@@ -437,13 +466,23 @@
 func (b *baseModuleContext) VisitDirectDepsWithTag(tag blueprint.DependencyTag, visit func(Module)) {
 	b.bp.VisitDirectDeps(func(module blueprint.Module) {
 		if b.bp.OtherModuleDependencyTag(module) == tag {
-			if aModule := b.validateAndroidModule(module, b.bp.OtherModuleDependencyTag(module), b.strictVisitDeps); aModule != nil {
+			if aModule := b.validateAndroidModule(module, tag, b.strictVisitDeps); aModule != nil {
 				visit(aModule)
 			}
 		}
 	})
 }
 
+func (b *baseModuleContext) VisitDirectDepsProxyWithTag(tag blueprint.DependencyTag, visit func(proxy ModuleProxy)) {
+	b.bp.VisitDirectDepsProxy(func(module blueprint.ModuleProxy) {
+		if b.bp.OtherModuleDependencyTag(module) == tag {
+			if aModule := b.validateAndroidModuleProxy(module, tag, b.strictVisitDeps); aModule != nil {
+				visit(*aModule)
+			}
+		}
+	})
+}
+
 func (b *baseModuleContext) VisitDirectDepsIf(pred func(Module) bool, visit func(Module)) {
 	b.bp.VisitDirectDepsIf(
 		// pred
@@ -505,7 +544,7 @@
 	})
 }
 
-func (b *baseModuleContext) WalkDepsProxy(visit func(Module, Module) bool) {
+func (b *baseModuleContext) WalkDepsProxy(visit func(ModuleProxy, ModuleProxy) bool) {
 	b.walkPath = []Module{ModuleProxy{blueprint.CreateModuleProxy(b.Module())}}
 	b.tagPath = []blueprint.DependencyTag{}
 	b.bp.WalkDepsProxy(func(child, parent blueprint.ModuleProxy) bool {
diff --git a/android/prebuilt.go b/android/prebuilt.go
index 4f04d05..bfd5595 100644
--- a/android/prebuilt.go
+++ b/android/prebuilt.go
@@ -362,10 +362,10 @@
 // the right module. This function is only safe to call after all TransitionMutators
 // have run, e.g. in GenerateAndroidBuildActions.
 func PrebuiltGetPreferred(ctx BaseModuleContext, module Module) Module {
-	if !module.IsReplacedByPrebuilt() {
+	if !OtherModuleProviderOrDefault(ctx, module, CommonPropertiesProviderKey).ReplacedByPrebuilt {
 		return module
 	}
-	if IsModulePrebuilt(module) {
+	if _, ok := OtherModuleProvider(ctx, module, PrebuiltModuleProviderKey); ok {
 		// If we're given a prebuilt then assume there's no source module around.
 		return module
 	}
@@ -373,11 +373,11 @@
 	sourceModDepFound := false
 	var prebuiltMod Module
 
-	ctx.WalkDeps(func(child, parent Module) bool {
+	ctx.WalkDepsProxy(func(child, parent ModuleProxy) bool {
 		if prebuiltMod != nil {
 			return false
 		}
-		if parent == ctx.Module() {
+		if ctx.EqualModules(parent, ctx.Module()) {
 			// First level: Only recurse if the module is found as a direct dependency.
 			sourceModDepFound = child == module
 			return sourceModDepFound
diff --git a/android/proto.go b/android/proto.go
index 0d8e097..66faa20 100644
--- a/android/proto.go
+++ b/android/proto.go
@@ -74,14 +74,14 @@
 		flags = append(flags, JoinWithPrefix(rootProtoIncludeDirs.Strings(), "-I"))
 	}
 
-	ctx.VisitDirectDepsWithTag(ProtoPluginDepTag, func(dep Module) {
-		if hostTool, ok := dep.(HostToolProvider); !ok || !hostTool.HostToolPath().Valid() {
+	ctx.VisitDirectDepsProxyWithTag(ProtoPluginDepTag, func(dep ModuleProxy) {
+		if h, ok := OtherModuleProvider(ctx, dep, HostToolProviderKey); !ok || !h.HostToolPath.Valid() {
 			ctx.PropertyErrorf("proto.plugin", "module %q is not a host tool provider",
 				ctx.OtherModuleName(dep))
 		} else {
 			plugin := String(p.Proto.Plugin)
-			deps = append(deps, hostTool.HostToolPath().Path())
-			flags = append(flags, "--plugin=protoc-gen-"+plugin+"="+hostTool.HostToolPath().String())
+			deps = append(deps, h.HostToolPath.Path())
+			flags = append(flags, "--plugin=protoc-gen-"+plugin+"="+h.HostToolPath.String())
 		}
 	})
 
diff --git a/cc/sanitize.go b/cc/sanitize.go
index 85fdb02..498ef7c 100644
--- a/cc/sanitize.go
+++ b/cc/sanitize.go
@@ -1830,10 +1830,7 @@
 
 type sanitizerLibraryDependencyTag struct {
 	blueprint.BaseDependencyTag
-}
-
-func (t sanitizerLibraryDependencyTag) AllowDisabledModuleDependency(target android.Module) bool {
-	return true
+	android.AlwaysAllowDisabledModuleDependencyTag
 }
 
 var _ android.AllowDisabledModuleDependency = (*sanitizerLibraryDependencyTag)(nil)
diff --git a/dexpreopt/config.go b/dexpreopt/config.go
index 84d4f10..3dd6f9a 100644
--- a/dexpreopt/config.go
+++ b/dexpreopt/config.go
@@ -462,6 +462,12 @@
 	return target.IsReplacedByPrebuilt()
 }
 
+func (d dex2oatDependencyTag) AllowDisabledModuleDependencyProxy(
+	ctx android.OtherModuleProviderContext, target android.ModuleProxy) bool {
+	return android.OtherModuleProviderOrDefault(
+		ctx, target, android.CommonPropertiesProviderKey).ReplacedByPrebuilt
+}
+
 // Dex2oatDepTag represents the dependency onto the dex2oatd module. It is added to any module that
 // needs dexpreopting and so it makes no sense for it to be checked for visibility or included in
 // the apex.
diff --git a/genrule/genrule.go b/genrule/genrule.go
index e5222a4..6eefb45 100644
--- a/genrule/genrule.go
+++ b/genrule/genrule.go
@@ -112,6 +112,12 @@
 	return target.IsReplacedByPrebuilt()
 }
 
+func (t hostToolDependencyTag) AllowDisabledModuleDependencyProxy(
+	ctx android.OtherModuleProviderContext, target android.ModuleProxy) bool {
+	return android.OtherModuleProviderOrDefault(
+		ctx, target, android.CommonPropertiesProviderKey).ReplacedByPrebuilt
+}
+
 var _ android.AllowDisabledModuleDependency = (*hostToolDependencyTag)(nil)
 
 type generatorProperties struct {
@@ -316,21 +322,18 @@
 	var packagedTools []android.PackagingSpec
 	if len(g.properties.Tools) > 0 {
 		seenTools := make(map[string]bool)
-
-		ctx.VisitDirectDepsAllowDisabled(func(module android.Module) {
-			switch tag := ctx.OtherModuleDependencyTag(module).(type) {
+		ctx.VisitDirectDepsProxyAllowDisabled(func(proxy android.ModuleProxy) {
+			switch tag := ctx.OtherModuleDependencyTag(proxy).(type) {
 			case hostToolDependencyTag:
-				tool := ctx.OtherModuleName(module)
 				// Necessary to retrieve any prebuilt replacement for the tool, since
 				// toolDepsMutator runs too late for the prebuilt mutators to have
 				// replaced the dependency.
-				module = android.PrebuiltGetPreferred(ctx, module)
-
-				switch t := module.(type) {
-				case android.HostToolProvider:
+				module := android.PrebuiltGetPreferred(ctx, proxy)
+				tool := ctx.OtherModuleName(module)
+				if h, ok := android.OtherModuleProvider(ctx, module, android.HostToolProviderKey); ok {
 					// A HostToolProvider provides the path to a tool, which will be copied
 					// into the sandbox.
-					if !t.(android.Module).Enabled(ctx) {
+					if !android.OtherModuleProviderOrDefault(ctx, module, android.CommonPropertiesProviderKey).Enabled {
 						if ctx.Config().AllowMissingDependencies() {
 							ctx.AddMissingDependencies([]string{tool})
 						} else {
@@ -338,13 +341,13 @@
 						}
 						return
 					}
-					path := t.HostToolPath()
+					path := h.HostToolPath
 					if !path.Valid() {
 						ctx.ModuleErrorf("host tool %q missing output file", tool)
 						return
 					}
 					if specs := android.OtherModuleProviderOrDefault(
-						ctx, t, android.InstallFilesProvider).TransitivePackagingSpecs.ToList(); specs != nil {
+						ctx, module, android.InstallFilesProvider).TransitivePackagingSpecs.ToList(); specs != nil {
 						// If the HostToolProvider has PackgingSpecs, which are definitions of the
 						// required relative locations of the tool and its dependencies, use those
 						// instead.  They will be copied to those relative locations in the sbox
@@ -366,7 +369,7 @@
 						tools = append(tools, path.Path())
 						addLocationLabel(tag.label, toolLocation{android.Paths{path.Path()}})
 					}
-				default:
+				} else {
 					ctx.ModuleErrorf("%q is not a host tool provider", tool)
 					return
 				}