Collect kythe, obj and tidy files using ModuleProxy.

Bug: 358427516
Test: Manually compare the generated mk and ninja files.
Change-Id: Ibee605372b1fae48ee9fbd8872b4ea0f6b1d5a7d
diff --git a/android/init.go b/android/init.go
index 1ace344..d3a13d0 100644
--- a/android/init.go
+++ b/android/init.go
@@ -20,6 +20,7 @@
 	gob.Register(extraFilesZip{})
 	gob.Register(InstallPath{})
 	gob.Register(ModuleGenPath{})
+	gob.Register(ModuleObjPath{})
 	gob.Register(ModuleOutPath{})
 	gob.Register(OutputPath{})
 	gob.Register(PhonyPath{})
diff --git a/android/module.go b/android/module.go
index a918d6e..72c8b94 100644
--- a/android/module.go
+++ b/android/module.go
@@ -1822,6 +1822,8 @@
 	Enabled bool
 	// Whether the module has been replaced by a prebuilt
 	ReplacedByPrebuilt bool
+	// The Target of artifacts that this module variant is responsible for creating.
+	CompileTarget Target
 }
 
 var CommonPropertiesProviderKey = blueprint.NewProvider[CommonPropertiesProviderData]()
@@ -2086,6 +2088,7 @@
 
 	commonData := CommonPropertiesProviderData{
 		ReplacedByPrebuilt: m.commonProperties.ReplacedByPrebuilt,
+		CompileTarget:      m.commonProperties.CompileTarget,
 	}
 	if m.commonProperties.ForcedDisabled {
 		commonData.Enabled = false
diff --git a/android/singleton.go b/android/singleton.go
index 913bf6a..6438182 100644
--- a/android/singleton.go
+++ b/android/singleton.go
@@ -64,6 +64,7 @@
 
 	VisitAllModulesBlueprint(visit func(blueprint.Module))
 	VisitAllModules(visit func(Module))
+	VisitAllModuleProxies(visit func(proxy ModuleProxy))
 	VisitAllModulesIf(pred func(Module) bool, visit func(Module))
 
 	VisitDirectDeps(module Module, visit func(Module))
@@ -77,6 +78,8 @@
 
 	VisitAllModuleVariants(module Module, visit func(Module))
 
+	VisitAllModuleVariantProxies(module Module, visit func(proxy ModuleProxy))
+
 	PrimaryModule(module Module) Module
 	FinalModule(module Module) Module
 
@@ -193,7 +196,7 @@
 }
 
 // visitAdaptor wraps a visit function that takes an android.Module parameter into
-// a function that takes an blueprint.Module parameter and only calls the visit function if the
+// a function that takes a blueprint.Module parameter and only calls the visit function if the
 // blueprint.Module is an android.Module.
 func visitAdaptor(visit func(Module)) func(blueprint.Module) {
 	return func(module blueprint.Module) {
@@ -203,6 +206,16 @@
 	}
 }
 
+// visitProxyAdaptor wraps a visit function that takes an android.ModuleProxy parameter into
+// a function that takes a blueprint.ModuleProxy parameter.
+func visitProxyAdaptor(visit func(proxy ModuleProxy)) func(proxy blueprint.ModuleProxy) {
+	return func(module blueprint.ModuleProxy) {
+		visit(ModuleProxy{
+			module: module,
+		})
+	}
+}
+
 // predAdaptor wraps a pred function that takes an android.Module parameter
 // into a function that takes an blueprint.Module parameter and only calls the visit function if the
 // blueprint.Module is an android.Module, otherwise returns false.
@@ -224,6 +237,10 @@
 	s.SingletonContext.VisitAllModules(visitAdaptor(visit))
 }
 
+func (s *singletonContextAdaptor) VisitAllModuleProxies(visit func(proxy ModuleProxy)) {
+	s.SingletonContext.VisitAllModuleProxies(visitProxyAdaptor(visit))
+}
+
 func (s *singletonContextAdaptor) VisitAllModulesIf(pred func(Module) bool, visit func(Module)) {
 	s.SingletonContext.VisitAllModulesIf(predAdaptor(pred), visitAdaptor(visit))
 }
@@ -248,6 +265,10 @@
 	s.SingletonContext.VisitAllModuleVariants(module, visitAdaptor(visit))
 }
 
+func (s *singletonContextAdaptor) VisitAllModuleVariantProxies(module Module, visit func(proxy ModuleProxy)) {
+	s.SingletonContext.VisitAllModuleVariantProxies(module, visitProxyAdaptor(visit))
+}
+
 func (s *singletonContextAdaptor) PrimaryModule(module Module) Module {
 	return s.SingletonContext.PrimaryModule(module).(Module)
 }
diff --git a/cc/cc.go b/cc/cc.go
index 5dee32e..069a861 100644
--- a/cc/cc.go
+++ b/cc/cc.go
@@ -43,6 +43,14 @@
 
 var CcMakeVarsInfoProvider = blueprint.NewProvider[*CcMakeVarsInfo]()
 
+type CcObjectInfo struct {
+	objFiles   android.Paths
+	tidyFiles  android.Paths
+	kytheFiles android.Paths
+}
+
+var CcObjectInfoProvider = blueprint.NewProvider[CcObjectInfo]()
+
 func init() {
 	RegisterCCBuildComponents(android.InitRegistrationContext)
 
@@ -648,10 +656,6 @@
 	installInRoot() bool
 }
 
-type xref interface {
-	XrefCcFiles() android.Paths
-}
-
 type overridable interface {
 	overriddenModules() []string
 }
@@ -900,12 +904,6 @@
 	staticAnalogue *StaticLibraryInfo
 
 	makeLinkType string
-	// Kythe (source file indexer) paths for this compilation module
-	kytheFiles android.Paths
-	// Object .o file output paths for this compilation module
-	objFiles android.Paths
-	// Tidy .tidy file output paths for this compilation module
-	tidyFiles android.Paths
 
 	// For apex variants, this is set as apex.min_sdk_version
 	apexSdkVersion android.ApiLevel
@@ -1473,10 +1471,6 @@
 	return isBionic(name)
 }
 
-func (c *Module) XrefCcFiles() android.Paths {
-	return c.kytheFiles
-}
-
 func (c *Module) isCfiAssemblySupportEnabled() bool {
 	return c.sanitize != nil &&
 		Bool(c.sanitize.Properties.Sanitize.Config.Cfi_assembly_support)
@@ -2048,9 +2042,6 @@
 		if ctx.Failed() {
 			return
 		}
-		c.kytheFiles = objs.kytheFiles
-		c.objFiles = objs.objFiles
-		c.tidyFiles = objs.tidyFiles
 	}
 
 	if c.linker != nil {
@@ -2115,6 +2106,17 @@
 		c.hasYacc = b.hasSrcExt(ctx, ".y") || b.hasSrcExt(ctx, ".yy")
 	}
 
+	ccObjectInfo := CcObjectInfo{
+		kytheFiles: objs.kytheFiles,
+	}
+	if !ctx.Config().KatiEnabled() || !android.ShouldSkipAndroidMkProcessing(ctx, c) {
+		ccObjectInfo.objFiles = objs.objFiles
+		ccObjectInfo.tidyFiles = objs.tidyFiles
+	}
+	if len(ccObjectInfo.kytheFiles)+len(ccObjectInfo.objFiles)+len(ccObjectInfo.tidyFiles) > 0 {
+		android.SetProvider(ctx, CcObjectInfoProvider, ccObjectInfo)
+	}
+
 	c.setOutputFiles(ctx)
 
 	if c.makeVarsInfo != nil {
@@ -2122,6 +2124,12 @@
 	}
 }
 
+func setOutputFilesIfNotEmpty(ctx ModuleContext, files android.Paths, tag string) {
+	if len(files) > 0 {
+		ctx.SetOutputFiles(files, tag)
+	}
+}
+
 func (c *Module) setOutputFiles(ctx ModuleContext) {
 	if c.outputFile.Valid() {
 		ctx.SetOutputFiles(android.Paths{c.outputFile.Path()}, "")
@@ -3958,9 +3966,10 @@
 
 func (ks *kytheExtractAllSingleton) GenerateBuildActions(ctx android.SingletonContext) {
 	var xrefTargets android.Paths
-	ctx.VisitAllModules(func(module android.Module) {
-		if ccModule, ok := module.(xref); ok {
-			xrefTargets = append(xrefTargets, ccModule.XrefCcFiles()...)
+	ctx.VisitAllModuleProxies(func(module android.ModuleProxy) {
+		files := android.OtherModuleProviderOrDefault(ctx, module, CcObjectInfoProvider).kytheFiles
+		if len(files) > 0 {
+			xrefTargets = append(xrefTargets, files...)
 		}
 	})
 	// TODO(asmundak): Perhaps emit a rule to output a warning if there were no xrefTargets
diff --git a/cc/tidy.go b/cc/tidy.go
index ec1e8a2..89bae17 100644
--- a/cc/tidy.go
+++ b/cc/tidy.go
@@ -219,15 +219,11 @@
 	subsetTidyFileGroups := make(map[string]android.Paths) // subset group name => tidy file Paths
 
 	// (1) Collect all obj/tidy files into OS-specific groups.
-	ctx.VisitAllModuleVariants(module, func(variant android.Module) {
-		if ctx.Config().KatiEnabled() && android.ShouldSkipAndroidMkProcessing(ctx, variant) {
-			return
-		}
-		if m, ok := variant.(*Module); ok {
-			osName := variant.Target().Os.Name
-			addToOSGroup(osName, m.objFiles, allObjFileGroups, subsetObjFileGroups)
-			addToOSGroup(osName, m.tidyFiles, allTidyFileGroups, subsetTidyFileGroups)
-		}
+	ctx.VisitAllModuleVariantProxies(module, func(variant android.ModuleProxy) {
+		osName := android.OtherModuleProviderOrDefault(ctx, variant, android.CommonPropertiesProviderKey).CompileTarget.Os.Name
+		info := android.OtherModuleProviderOrDefault(ctx, variant, CcObjectInfoProvider)
+		addToOSGroup(osName, info.objFiles, allObjFileGroups, subsetObjFileGroups)
+		addToOSGroup(osName, info.tidyFiles, allTidyFileGroups, subsetTidyFileGroups)
 	})
 
 	// (2) Add an all-OS group, with "" or "subset" name, to include all os-specific phony targets.