Add GetDirectDepProxyWithTag with some refactoring work.

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

2 files changed, 43 insertions, 21 deletions

diff --git a/android/base_module_context.go b/android/base_module_context.go
index 839afaa66..d8558d0e5 100644
--- a/android/base_module_context.go
+++ b/android/base_module_context.go
@@ -117,6 +117,8 @@ type BaseModuleContext interface {
 	// dependencies that are not an android.Module.
 	GetDirectDepWithTag(name string, tag blueprint.DependencyTag) Module
 
+	GetDirectDepProxyWithTag(name string, tag blueprint.DependencyTag) *ModuleProxy
+
 	// VisitDirectDeps 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.  It raises an error if any of the
@@ -311,6 +313,13 @@ func (b *baseModuleContext) GetDirectDepWithTag(name string, tag blueprint.Depen
 	return nil
 }
 
+func (b *baseModuleContext) GetDirectDepProxyWithTag(name string, tag blueprint.DependencyTag) *ModuleProxy {
+	if module := b.bp.GetDirectDepProxyWithTag(name, tag); module != nil {
+		return &ModuleProxy{*module}
+	}
+	return nil
+}
+
 func (b *baseModuleContext) blueprintBaseModuleContext() blueprint.BaseModuleContext {
 	return b.bp
 }
@@ -399,34 +408,30 @@ func (b *baseModuleContext) validateAndroidModuleProxy(
 	return &aModule
 }
 
-type dep struct {
-	mod blueprint.Module
-	tag blueprint.DependencyTag
-}
-
-func (b *baseModuleContext) getDirectDepsInternal(name string, tag blueprint.DependencyTag) []dep {
-	var deps []dep
+func (b *baseModuleContext) getDirectDepsInternal(name string, tag blueprint.DependencyTag) []Module {
+	var deps []Module
 	b.VisitDirectDeps(func(module Module) {
 		if module.base().BaseModuleName() == name {
 			returnedTag := b.bp.OtherModuleDependencyTag(module)
 			if tag == nil || returnedTag == tag {
-				deps = append(deps, dep{module, returnedTag})
+				deps = append(deps, module)
 			}
 		}
 	})
 	return deps
 }
 
-func (b *baseModuleContext) getDirectDepInternal(name string, tag blueprint.DependencyTag) (blueprint.Module, blueprint.DependencyTag) {
-	deps := b.getDirectDepsInternal(name, tag)
-	if len(deps) == 1 {
-		return deps[0].mod, deps[0].tag
-	} else if len(deps) >= 2 {
-		panic(fmt.Errorf("Multiple dependencies having same BaseModuleName() %q found from %q",
-			name, b.ModuleName()))
-	} else {
-		return nil, nil
-	}
+func (b *baseModuleContext) getDirectDepsProxyInternal(name string, tag blueprint.DependencyTag) []ModuleProxy {
+	var deps []ModuleProxy
+	b.VisitDirectDepsProxy(func(module ModuleProxy) {
+		if OtherModuleProviderOrDefault(b, module, CommonModuleInfoKey).BaseModuleName == name {
+			returnedTag := b.OtherModuleDependencyTag(module)
+			if tag == nil || returnedTag == tag {
+				deps = append(deps, module)
+			}
+		}
+	})
+	return deps
 }
 
 func (b *baseModuleContext) GetDirectDepsWithTag(tag blueprint.DependencyTag) []Module {
diff --git a/android/module_context.go b/android/module_context.go
index b097117d8..b59e33d83 100644
--- a/android/module_context.go
+++ b/android/module_context.go
@@ -440,10 +440,27 @@ func (m *moduleContext) GetMissingDependencies() []string {
 }
 
 func (m *moduleContext) GetDirectDepWithTag(name string, tag blueprint.DependencyTag) Module {
-	if module, _ := m.getDirectDepInternal(name, tag); module != nil {
-		return module.(Module)
+	deps := m.getDirectDepsInternal(name, tag)
+	if len(deps) == 1 {
+		return deps[0]
+	} else if len(deps) >= 2 {
+		panic(fmt.Errorf("Multiple dependencies having same BaseModuleName() %q found from %q",
+			name, m.ModuleName()))
+	} else {
+		return nil
+	}
+}
+
+func (m *moduleContext) GetDirectDepProxyWithTag(name string, tag blueprint.DependencyTag) *ModuleProxy {
+	deps := m.getDirectDepsProxyInternal(name, tag)
+	if len(deps) == 1 {
+		return &deps[0]
+	} else if len(deps) >= 2 {
+		panic(fmt.Errorf("Multiple dependencies having same BaseModuleName() %q found from %q",
+			name, m.ModuleName()))
+	} else {
+		return nil
 	}
-	return nil
 }
 
 func (m *moduleContext) ModuleSubDir() string {