Change GetModuleFromPathDep to use ModuleProxy.

Bug: 377723687
Test: Compare ninja and mk files generated.
Change-Id: I428b0965b217adb20a792ebde88374e0c6fae9d6

9 files changed, 107 insertions, 53 deletions

diff --git a/android/base_module_context.go b/android/base_module_context.go
index 060fae5bc..719d6d5bf 100644
--- a/android/base_module_context.go
+++ b/android/base_module_context.go
@@ -113,7 +113,7 @@ type BaseModuleContext interface {
 	// GetDirectDepWithTag returns the Module the direct dependency with the specified name, or nil if
 	// none exists.  It panics if the dependency does not have the specified tag.  It skips any
 	// dependencies that are not an android.Module.
-	GetDirectDepWithTag(name string, tag blueprint.DependencyTag) blueprint.Module
+	GetDirectDepWithTag(name string, tag blueprint.DependencyTag) Module
 
 	// GetDirectDep returns the Module and DependencyTag for the direct dependency with the specified
 	// name, or nil if none exists.  If there are multiple dependencies on the same module it returns
@@ -129,13 +129,14 @@ type BaseModuleContext interface {
 	// function, it may be invalidated by future mutators.
 	VisitDirectDeps(visit func(Module))
 
-	// VisitDirectDeps calls visit for each direct dependency.  If there are multiple
+	// VisitDirectDepsProxy calls visit for each direct dependency.  If there are 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.
+	// and OtherModuleDependencyTag will return a different tag for each. It raises an error if any of the
+	// dependencies are disabled.
 	//
-	// The Module passed to the visit function should not be retained outside of the visit
+	// The ModuleProxy passed to the visit function should not be retained outside of the visit
 	// function, it may be invalidated by future mutators.
-	VisitDirectDepsAllowDisabled(visit func(Module))
+	VisitDirectDepsProxy(visit func(proxy ModuleProxy))
 
 	// VisitDirectDepsProxyAllowDisabled calls visit for each direct dependency.  If there are
 	// multiple direct dependencies on the same module visit will be called multiple times on
@@ -261,7 +262,9 @@ func (b *baseModuleContext) EqualModules(m1, m2 Module) bool {
 func (b *baseModuleContext) OtherModuleName(m blueprint.Module) string {
 	return b.bp.OtherModuleName(getWrappedModule(m))
 }
-func (b *baseModuleContext) OtherModuleDir(m blueprint.Module) string { return b.bp.OtherModuleDir(m) }
+func (b *baseModuleContext) OtherModuleDir(m blueprint.Module) string {
+	return b.bp.OtherModuleDir(getWrappedModule(m))
+}
 func (b *baseModuleContext) OtherModuleErrorf(m blueprint.Module, fmt string, args ...interface{}) {
 	b.bp.OtherModuleErrorf(m, fmt, args...)
 }
@@ -298,8 +301,11 @@ func (b *baseModuleContext) setProvider(provider blueprint.AnyProviderKey, value
 	b.bp.SetProvider(provider, value)
 }
 
-func (b *baseModuleContext) GetDirectDepWithTag(name string, tag blueprint.DependencyTag) blueprint.Module {
-	return b.bp.GetDirectDepWithTag(name, tag)
+func (b *baseModuleContext) GetDirectDepWithTag(name string, tag blueprint.DependencyTag) Module {
+	if module := b.bp.GetDirectDepWithTag(name, tag); module != nil {
+		return module.(Module)
+	}
+	return nil
 }
 
 func (b *baseModuleContext) blueprintBaseModuleContext() blueprint.BaseModuleContext {
@@ -464,18 +470,16 @@ func (b *baseModuleContext) VisitDirectDeps(visit func(Module)) {
 	})
 }
 
-func (b *baseModuleContext) VisitDirectDepsAllowDisabled(visit func(Module)) {
-	b.bp.VisitDirectDeps(func(module blueprint.Module) {
-		visit(module.(Module))
+func (b *baseModuleContext) VisitDirectDepsProxy(visit func(ModuleProxy)) {
+	b.bp.VisitDirectDepsProxy(func(module blueprint.ModuleProxy) {
+		if aModule := b.validateAndroidModuleProxy(module, b.bp.OtherModuleDependencyTag(module), b.strictVisitDeps); aModule != nil {
+			visit(*aModule)
+		}
 	})
 }
 
 func (b *baseModuleContext) VisitDirectDepsProxyAllowDisabled(visit func(proxy ModuleProxy)) {
-	b.bp.VisitDirectDepsProxy(func(module blueprint.ModuleProxy) {
-		visit(ModuleProxy{
-			module: module,
-		})
-	})
+	b.bp.VisitDirectDepsProxy(visitProxyAdaptor(visit))
 }
 
 func (b *baseModuleContext) VisitDirectDepsWithTag(tag blueprint.DependencyTag, visit func(Module)) {
diff --git a/android/module.go b/android/module.go
index 3b30c11ef..fba989f69 100644
--- a/android/module.go
+++ b/android/module.go
@@ -1834,6 +1834,12 @@ type InstallFilesInfo struct {
 
 var InstallFilesProvider = blueprint.NewProvider[InstallFilesInfo]()
 
+type SourceFilesInfo struct {
+	Srcs Paths
+}
+
+var SourceFilesInfoKey = blueprint.NewProvider[SourceFilesInfo]()
+
 type FinalModuleBuildTargetsInfo struct {
 	// Used by buildTargetSingleton to create checkbuild and per-directory build targets
 	// Only set on the final variant of each module
@@ -2038,6 +2044,10 @@ func (m *ModuleBase) GenerateBuildActions(blueprintCtx blueprint.ModuleContext)
 		ctx.GetMissingDependencies()
 	}
 
+	if sourceFileProducer, ok := m.module.(SourceFileProducer); ok {
+		SetProvider(ctx, SourceFilesInfoKey, SourceFilesInfo{Srcs: sourceFileProducer.Srcs()})
+	}
+
 	if ctx.IsFinalModule(m.module) {
 		m.generateModuleTarget(ctx)
 		if ctx.Failed() {
@@ -2634,7 +2644,7 @@ type SourceFileProducer interface {
 
 // OutputFilesForModule returns the output file paths with the given tag. On error, including if the
 // module produced zero paths, it reports errors to the ctx and returns nil.
-func OutputFilesForModule(ctx PathContext, module blueprint.Module, tag string) Paths {
+func OutputFilesForModule(ctx PathContext, module Module, tag string) Paths {
 	paths, err := outputFilesForModule(ctx, module, tag)
 	if err != nil {
 		reportPathError(ctx, err)
@@ -2645,7 +2655,7 @@ func OutputFilesForModule(ctx PathContext, module blueprint.Module, tag string)
 
 // OutputFileForModule returns the output file paths with the given tag.  On error, including if the
 // module produced zero or multiple paths, it reports errors to the ctx and returns nil.
-func OutputFileForModule(ctx PathContext, module blueprint.Module, tag string) Path {
+func OutputFileForModule(ctx PathContext, module Module, tag string) Path {
 	paths, err := outputFilesForModule(ctx, module, tag)
 	if err != nil {
 		reportPathError(ctx, err)
@@ -2678,20 +2688,34 @@ func OutputFileForModule(ctx PathContext, module blueprint.Module, tag string) P
 	return paths[0]
 }
 
-func outputFilesForModule(ctx PathContext, module blueprint.Module, tag string) (Paths, error) {
+type OutputFilesProviderModuleContext interface {
+	OtherModuleProviderContext
+	Module() Module
+	GetOutputFiles() OutputFilesInfo
+	EqualModules(m1, m2 Module) bool
+}
+
+func outputFilesForModule(ctx PathContext, module Module, tag string) (Paths, error) {
 	outputFilesFromProvider, err := outputFilesForModuleFromProvider(ctx, module, tag)
 	if outputFilesFromProvider != nil || err != OutputFilesProviderNotSet {
 		return outputFilesFromProvider, err
 	}
-	if sourceFileProducer, ok := module.(SourceFileProducer); ok {
-		if tag != "" {
-			return nil, fmt.Errorf("module %q is a SourceFileProducer, which does not support tag %q", pathContextName(ctx, module), tag)
+
+	if octx, ok := ctx.(OutputFilesProviderModuleContext); ok {
+		if octx.EqualModules(octx.Module(), module) {
+			if sourceFileProducer, ok := module.(SourceFileProducer); ok {
+				return sourceFileProducer.Srcs(), nil
+			}
+		} else if sourceFiles, ok := OtherModuleProvider(octx, module, SourceFilesInfoKey); ok {
+			if tag != "" {
+				return nil, fmt.Errorf("module %q is a SourceFileProducer, which does not support tag %q", pathContextName(ctx, module), tag)
+			}
+			paths := sourceFiles.Srcs
+			return paths, nil
 		}
-		paths := sourceFileProducer.Srcs()
-		return paths, nil
-	} else {
-		return nil, fmt.Errorf("module %q is not a SourceFileProducer or having valid output file for tag %q", pathContextName(ctx, module), tag)
 	}
+
+	return nil, fmt.Errorf("module %q is not a SourceFileProducer or having valid output file for tag %q", pathContextName(ctx, module), tag)
 }
 
 // This method uses OutputFilesProvider for output files
@@ -2700,26 +2724,19 @@ func outputFilesForModule(ctx PathContext, module blueprint.Module, tag string)
 // from outputFiles property of module base, to avoid both setting and
 // reading OutputFilesProvider before GenerateBuildActions is finished.
 // If a module doesn't have the OutputFilesProvider, nil is returned.
-func outputFilesForModuleFromProvider(ctx PathContext, module blueprint.Module, tag string) (Paths, error) {
+func outputFilesForModuleFromProvider(ctx PathContext, module Module, tag string) (Paths, error) {
 	var outputFiles OutputFilesInfo
 	fromProperty := false
 
-	type OutputFilesProviderModuleContext interface {
-		OtherModuleProviderContext
-		Module() Module
-		GetOutputFiles() OutputFilesInfo
-	}
-
 	if mctx, isMctx := ctx.(OutputFilesProviderModuleContext); isMctx {
-		if mctx.Module() != module {
+		if !mctx.EqualModules(mctx.Module(), module) {
 			outputFiles, _ = OtherModuleProvider(mctx, module, OutputFilesProvider)
 		} else {
 			outputFiles = mctx.GetOutputFiles()
 			fromProperty = true
 		}
 	} else if cta, isCta := ctx.(*singletonContextAdaptor); isCta {
-		providerData, _ := cta.otherModuleProvider(module, OutputFilesProvider)
-		outputFiles, _ = providerData.(OutputFilesInfo)
+		outputFiles, _ = OtherModuleProvider(cta, module, OutputFilesProvider)
 	} else {
 		return nil, fmt.Errorf("unsupported context %q in method outputFilesForModuleFromProvider", reflect.TypeOf(ctx))
 	}
diff --git a/android/module_context.go b/android/module_context.go
index 41cb0ccb2..20149074e 100644
--- a/android/module_context.go
+++ b/android/module_context.go
@@ -16,13 +16,13 @@ package android
 
 import (
 	"fmt"
-	"github.com/google/blueprint/depset"
 	"path"
 	"path/filepath"
 	"slices"
 	"strings"
 
 	"github.com/google/blueprint"
+	"github.com/google/blueprint/depset"
 	"github.com/google/blueprint/proptools"
 )
 
@@ -439,9 +439,11 @@ func (m *moduleContext) GetMissingDependencies() []string {
 	return missingDeps
 }
 
-func (m *moduleContext) GetDirectDepWithTag(name string, tag blueprint.DependencyTag) blueprint.Module {
-	module, _ := m.getDirectDepInternal(name, tag)
-	return module
+func (m *moduleContext) GetDirectDepWithTag(name string, tag blueprint.DependencyTag) Module {
+	if module, _ := m.getDirectDepInternal(name, tag); module != nil {
+		return module.(Module)
+	}
+	return nil
 }
 
 func (m *moduleContext) ModuleSubDir() string {
diff --git a/android/module_test.go b/android/module_test.go
index d76d9b329..d5bf94137 100644
--- a/android/module_test.go
+++ b/android/module_test.go
@@ -998,6 +998,10 @@ func (p *pathContextAddMissingDependenciesWrapper) GetOutputFiles() OutputFilesI
 	return OutputFilesInfo{}
 }
 
+func (p *pathContextAddMissingDependenciesWrapper) EqualModules(m1, m2 Module) bool {
+	return m1 == m2
+}
+
 func TestOutputFileForModule(t *testing.T) {
 	testcases := []struct {
 		name        string
diff --git a/android/path_properties_test.go b/android/path_properties_test.go
index 07b48696c..6f44f2872 100644
--- a/android/path_properties_test.go
+++ b/android/path_properties_test.go
@@ -64,7 +64,7 @@ func (p *pathDepsMutatorTestModule) GenerateAndroidBuildActions(ctx ModuleContex
 	if p.props.Foo != "" {
 		// Make sure there is only one dependency on a module listed in a property present in multiple property structs
 		m := SrcIsModule(p.props.Foo)
-		if GetModuleFromPathDep(ctx, m, "") == nil {
+		if GetModuleProxyFromPathDep(ctx, m, "") == nil {
 			ctx.ModuleErrorf("GetDirectDepWithTag failed")
 		}
 	}
diff --git a/android/paths.go b/android/paths.go
index 9cb872d6f..3f67c55b2 100644
--- a/android/paths.go
+++ b/android/paths.go
@@ -91,6 +91,8 @@ type ModuleWithDepsPathContext interface {
 	EarlyModulePathContext
 	OtherModuleProviderContext
 	VisitDirectDeps(visit func(Module))
+	VisitDirectDepsProxy(visit func(ModuleProxy))
+	VisitDirectDepsProxyWithTag(tag blueprint.DependencyTag, visit func(ModuleProxy))
 	OtherModuleDependencyTag(m blueprint.Module) blueprint.DependencyTag
 	HasMutatorFinished(mutatorName string) bool
 }
@@ -598,7 +600,7 @@ func DirectoryPathsForModuleSrc(ctx ModuleMissingDepsPathContext, paths []string
 
 	for _, path := range paths {
 		if m, t := SrcIsModuleWithTag(path); m != "" {
-			module := GetModuleFromPathDep(ctx, m, t)
+			module := GetModuleProxyFromPathDep(ctx, m, t)
 			if module == nil {
 				ctx.ModuleErrorf(`missing dependency on %q, is the property annotated with android:"path"?`, m)
 				continue
@@ -611,7 +613,7 @@ func DirectoryPathsForModuleSrc(ctx ModuleMissingDepsPathContext, paths []string
 			if !ok {
 				panic(fmt.Errorf("%s is not an OtherModuleProviderContext", ctx))
 			}
-			if dirProvider, ok := OtherModuleProvider(mctx, module, DirProvider); ok {
+			if dirProvider, ok := OtherModuleProvider(mctx, *module, DirProvider); ok {
 				ret = append(ret, dirProvider.Dirs...)
 			} else {
 				ReportPathErrorf(ctx, "module %q does not implement DirProvider", module)
@@ -669,14 +671,15 @@ func (p OutputPaths) Strings() []string {
 // If the dependency is not found, a missingErrorDependency is returned.
 // If the module dependency is not a SourceFileProducer or OutputFileProducer, appropriate errors will be returned.
 func getPathsFromModuleDep(ctx ModuleWithDepsPathContext, path, moduleName, tag string) (Paths, error) {
-	module := GetModuleFromPathDep(ctx, moduleName, tag)
+	module := GetModuleProxyFromPathDep(ctx, moduleName, tag)
 	if module == nil {
 		return nil, missingDependencyError{[]string{moduleName}}
 	}
-	if aModule, ok := module.(Module); ok && !aModule.Enabled(ctx) {
+	if !OtherModuleProviderOrDefault(ctx, *module, CommonPropertiesProviderKey).Enabled {
 		return nil, missingDependencyError{[]string{moduleName}}
 	}
-	outputFiles, err := outputFilesForModule(ctx, module, tag)
+
+	outputFiles, err := outputFilesForModule(ctx, *module, tag)
 	if outputFiles != nil && err == nil {
 		return outputFiles, nil
 	} else {
@@ -684,7 +687,7 @@ func getPathsFromModuleDep(ctx ModuleWithDepsPathContext, path, moduleName, tag
 	}
 }
 
-// GetModuleFromPathDep will return the module that was added as a dependency automatically for
+// GetModuleProxyFromPathDep will return the module that was added as a dependency automatically for
 // properties tagged with `android:"path"` or manually using ExtractSourceDeps or
 // ExtractSourcesDeps.
 //
@@ -694,6 +697,27 @@ func getPathsFromModuleDep(ctx ModuleWithDepsPathContext, path, moduleName, tag
 //
 // If tag is "" then the returned module will be the dependency that was added for ":moduleName".
 // Otherwise, it is the dependency that was added for ":moduleName{tag}".
+func GetModuleProxyFromPathDep(ctx ModuleWithDepsPathContext, moduleName, tag string) *ModuleProxy {
+	var found *ModuleProxy
+	// The sourceOrOutputDepTag uniquely identifies the module dependency as it contains both the
+	// module name and the tag. Dependencies added automatically for properties tagged with
+	// `android:"path"` are deduped so are guaranteed to be unique. It is possible for duplicate
+	// dependencies to be added manually using ExtractSourcesDeps or ExtractSourceDeps but even then
+	// it will always be the case that the dependencies will be identical, i.e. the same tag and same
+	// moduleName referring to the same dependency module.
+	//
+	// It does not matter whether the moduleName is a fully qualified name or if the module
+	// dependency is a prebuilt module. All that matters is the same information is supplied to
+	// create the tag here as was supplied to create the tag when the dependency was added so that
+	// this finds the matching dependency module.
+	expectedTag := sourceOrOutputDepTag(moduleName, tag)
+	ctx.VisitDirectDepsProxyWithTag(expectedTag, func(module ModuleProxy) {
+		found = &module
+	})
+	return found
+}
+
+// Deprecated: use GetModuleProxyFromPathDep
 func GetModuleFromPathDep(ctx ModuleWithDepsPathContext, moduleName, tag string) blueprint.Module {
 	var found blueprint.Module
 	// The sourceOrOutputDepTag uniquely identifies the module dependency as it contains both the
diff --git a/apex/builder.go b/apex/builder.go
index 305d5092a..b04a9d729 100644
--- a/apex/builder.go
+++ b/apex/builder.go
@@ -397,8 +397,10 @@ func (a *apexBundle) buildFileContexts(ctx android.ModuleContext) android.Output
 	} else {
 		if m, t := android.SrcIsModuleWithTag(*a.properties.File_contexts); m != "" {
 			isFileContextsModule = true
-			otherModule := android.GetModuleFromPathDep(ctx, m, t)
-			fileContextsDir = ctx.OtherModuleDir(otherModule)
+			otherModule := android.GetModuleProxyFromPathDep(ctx, m, t)
+			if otherModule != nil {
+				fileContextsDir = ctx.OtherModuleDir(*otherModule)
+			}
 		}
 		fileContexts = android.PathForModuleSrc(ctx, *a.properties.File_contexts)
 	}
diff --git a/java/classpath_fragment.go b/java/classpath_fragment.go
index 18a5dae6c..fdccd3a84 100644
--- a/java/classpath_fragment.go
+++ b/java/classpath_fragment.go
@@ -18,9 +18,10 @@ package java
 
 import (
 	"fmt"
+	"strings"
+
 	"github.com/google/blueprint"
 	"github.com/google/blueprint/proptools"
-	"strings"
 
 	"android/soong/android"
 )
@@ -103,7 +104,7 @@ type classpathJar struct {
 func gatherPossibleApexModuleNamesAndStems(ctx android.ModuleContext, contents []string, tag blueprint.DependencyTag) []string {
 	set := map[string]struct{}{}
 	for _, name := range contents {
-		dep, _ := ctx.GetDirectDepWithTag(name, tag).(android.Module)
+		dep := ctx.GetDirectDepWithTag(name, tag)
 		set[ModuleStemForDeapexing(dep)] = struct{}{}
 		if m, ok := dep.(ModuleWithStem); ok {
 			set[m.Stem()] = struct{}{}
diff --git a/java/droiddoc.go b/java/droiddoc.go
index 2980d91de..82713920d 100644
--- a/java/droiddoc.go
+++ b/java/droiddoc.go
@@ -427,9 +427,9 @@ func (j *Javadoc) collectDeps(ctx android.ModuleContext) deps {
 	// Find the corresponding aconfig_declarations module name for such case.
 	for _, src := range j.properties.Srcs {
 		if moduleName, tag := android.SrcIsModuleWithTag(src); moduleName != "" {
-			otherModule := android.GetModuleFromPathDep(ctx, moduleName, tag)
+			otherModule := android.GetModuleProxyFromPathDep(ctx, moduleName, tag)
 			if otherModule != nil {
-				if dep, ok := android.OtherModuleProvider(ctx, otherModule, android.CodegenInfoProvider); ok {
+				if dep, ok := android.OtherModuleProvider(ctx, *otherModule, android.CodegenInfoProvider); ok {
 					deps.aconfigProtoFiles = append(deps.aconfigProtoFiles, dep.IntermediateCacheOutputPaths...)
 				}
 			}