bp2build include generated hdrs in all src splits

Currently generated headers are only included in the srcs split;
however, if c or assembly sources depend on the headers, the
dependendency is lacking and can cause non-deterministic failures.

Test: build/bazel/ci/bp2build.sh
Change-Id: Ic354c88300a5fb1b85241f68889e904c33a81488
diff --git a/bp2build/cc_library_static_conversion_test.go b/bp2build/cc_library_static_conversion_test.go
index 9f6f450..ef2ec32 100644
--- a/bp2build/cc_library_static_conversion_test.go
+++ b/bp2build/cc_library_static_conversion_test.go
@@ -986,6 +986,41 @@
 	})
 }
 
+func TestCcLibraryStaticGeneratedHeadersAllPartitions(t *testing.T) {
+	runCcLibraryStaticTestCase(t, bp2buildTestCase{
+		moduleTypeUnderTest:                "cc_library_static",
+		moduleTypeUnderTestFactory:         cc.LibraryStaticFactory,
+		moduleTypeUnderTestBp2BuildMutator: cc.CcLibraryStaticBp2Build,
+		blueprint: soongCcLibraryStaticPreamble + `
+genrule {
+    name: "generated_hdr",
+    cmd: "nothing to see here",
+}
+
+cc_library_static {
+    name: "foo_static",
+    srcs: ["cpp_src.cpp", "as_src.S", "c_src.c"],
+    generated_headers: ["generated_hdr"],
+    include_build_directory: false,
+}`,
+		expectedBazelTargets: []string{`cc_library_static(
+    name = "foo_static",
+    srcs = [
+        "cpp_src.cpp",
+        ":generated_hdr",
+    ],
+    srcs_as = [
+        "as_src.S",
+        ":generated_hdr",
+    ],
+    srcs_c = [
+        "c_src.c",
+        ":generated_hdr",
+    ],
+)`},
+	})
+}
+
 func TestCcLibraryStaticArchSrcsExcludeSrcsGeneratedFiles(t *testing.T) {
 	runCcLibraryStaticTestCase(t, bp2buildTestCase{
 		description:                        "cc_library_static arch srcs/exclude_srcs with generated files",
@@ -1067,10 +1102,10 @@
     name = "foo_static3",
     srcs = [
         "common.cpp",
-        ":generated_hdr",
-        "//dep:generated_hdr_other_pkg",
         ":generated_src",
         "//dep:generated_src_other_pkg",
+        ":generated_hdr",
+        "//dep:generated_hdr_other_pkg",
     ] + select({
         "//build/bazel/platforms/arch:x86": [
             "for-x86.cpp",
@@ -1082,8 +1117,8 @@
         ],
     }) + select({
         "//build/bazel/platforms/os:android": [
-            "//dep:generated_hdr_other_pkg_android",
             ":generated_src_android",
+            "//dep:generated_hdr_other_pkg_android",
         ],
         "//conditions:default": [],
     }),
diff --git a/cc/bp2build.go b/cc/bp2build.go
index c078096..a7cabcc 100644
--- a/cc/bp2build.go
+++ b/cc/bp2build.go
@@ -26,6 +26,12 @@
 	"github.com/google/blueprint/proptools"
 )
 
+const (
+	cSrcPartition   = "c"
+	asSrcPartition  = "as"
+	cppSrcPartition = "cpp"
+)
+
 // staticOrSharedAttributes are the Bazel-ified versions of StaticOrSharedProperties --
 // properties which apply to either the shared or static version of a cc_library module.
 type staticOrSharedAttributes struct {
@@ -43,7 +49,7 @@
 	System_dynamic_deps bazel.LabelListAttribute
 }
 
-func groupSrcsByExtension(ctx android.TopDownMutatorContext, srcs bazel.LabelListAttribute) (cppSrcs, cSrcs, asSrcs bazel.LabelListAttribute) {
+func groupSrcsByExtension(ctx android.TopDownMutatorContext, srcs bazel.LabelListAttribute) bazel.PartitionToLabelListAttribute {
 	// Check that a module is a filegroup type named <label>.
 	isFilegroupNamed := func(m android.Module, fullLabel string) bool {
 		if ctx.OtherModuleType(m) != "filegroup" {
@@ -75,17 +81,14 @@
 
 	// TODO(b/190006308): Handle language detection of sources in a Bazel rule.
 	partitioned := bazel.PartitionLabelListAttribute(ctx, &srcs, bazel.LabelPartitions{
-		"c":  bazel.LabelPartition{Extensions: []string{".c"}, LabelMapper: addSuffixForFilegroup("_c_srcs")},
-		"as": bazel.LabelPartition{Extensions: []string{".s", ".S"}, LabelMapper: addSuffixForFilegroup("_as_srcs")},
+		cSrcPartition:  bazel.LabelPartition{Extensions: []string{".c"}, LabelMapper: addSuffixForFilegroup("_c_srcs")},
+		asSrcPartition: bazel.LabelPartition{Extensions: []string{".s", ".S"}, LabelMapper: addSuffixForFilegroup("_as_srcs")},
 		// C++ is the "catch-all" group, and comprises generated sources because we don't
 		// know the language of these sources until the genrule is executed.
-		"cpp": bazel.LabelPartition{Extensions: []string{".cpp", ".cc", ".cxx", ".mm"}, LabelMapper: addSuffixForFilegroup("_cpp_srcs"), Keep_remainder: true},
+		cppSrcPartition: bazel.LabelPartition{Extensions: []string{".cpp", ".cc", ".cxx", ".mm"}, LabelMapper: addSuffixForFilegroup("_cpp_srcs"), Keep_remainder: true},
 	})
 
-	cSrcs = partitioned["c"]
-	asSrcs = partitioned["as"]
-	cppSrcs = partitioned["cpp"]
-	return
+	return partitioned
 }
 
 // bp2BuildParseLibProps returns the attributes for a variant of a cc_library.
@@ -175,10 +178,10 @@
 		}
 	}
 
-	cppSrcs, cSrcs, asSrcs := groupSrcsByExtension(ctx, attrs.Srcs)
-	attrs.Srcs = cppSrcs
-	attrs.Srcs_c = cSrcs
-	attrs.Srcs_as = asSrcs
+	partitionedSrcs := groupSrcsByExtension(ctx, attrs.Srcs)
+	attrs.Srcs = partitionedSrcs[cppSrcPartition]
+	attrs.Srcs_c = partitionedSrcs[cSrcPartition]
+	attrs.Srcs_as = partitionedSrcs[asSrcPartition]
 
 	return attrs
 }
@@ -239,6 +242,7 @@
 // bp2BuildParseCompilerProps returns copts, srcs and hdrs and other attributes.
 func bp2BuildParseCompilerProps(ctx android.TopDownMutatorContext, module *Module) compilerAttributes {
 	var srcs bazel.LabelListAttribute
+	var implementationHdrs bazel.LabelListAttribute
 	var copts bazel.StringListAttribute
 	var asFlags bazel.StringListAttribute
 	var conlyFlags bazel.StringListAttribute
@@ -265,10 +269,8 @@
 		anySrcs := false
 		// Add srcs-like dependencies such as generated files.
 		// First create a LabelList containing these dependencies, then merge the values with srcs.
-		generatedHdrsAndSrcs := props.Generated_headers
-		generatedHdrsAndSrcs = append(generatedHdrsAndSrcs, props.Generated_sources...)
-		generatedHdrsAndSrcsLabelList := android.BazelLabelForModuleDepsExcludes(ctx, generatedHdrsAndSrcs, props.Exclude_generated_sources)
-		if len(generatedHdrsAndSrcs) > 0 || len(props.Exclude_generated_sources) > 0 {
+		generatedSrcsLabelList := android.BazelLabelForModuleDepsExcludes(ctx, props.Generated_sources, props.Exclude_generated_sources)
+		if len(props.Generated_sources) > 0 || len(props.Exclude_generated_sources) > 0 {
 			anySrcs = true
 		}
 
@@ -276,7 +278,7 @@
 		if len(props.Srcs) > 0 || len(props.Exclude_srcs) > 0 {
 			anySrcs = true
 		}
-		return bazel.AppendBazelLabelLists(allSrcsLabelList, generatedHdrsAndSrcsLabelList), anySrcs
+		return bazel.AppendBazelLabelLists(allSrcsLabelList, generatedSrcsLabelList), anySrcs
 	}
 
 	archVariantCompilerProps := module.GetArchVariantProperties(ctx, &BaseCompilerProperties{})
@@ -288,6 +290,9 @@
 				if srcsList, ok := parseSrcs(baseCompilerProps); ok {
 					srcs.SetSelectValue(axis, config, srcsList)
 				}
+				if len(baseCompilerProps.Generated_headers) > 0 {
+					implementationHdrs.SetSelectValue(axis, config, android.BazelLabelForModuleDeps(ctx, baseCompilerProps.Generated_headers))
+				}
 
 				if axis == bazel.NoConfigAxis {
 					// If cpp_std is not specified, don't generate it in the
@@ -329,6 +334,21 @@
 	}
 
 	srcs.ResolveExcludes()
+	partitionedSrcs := groupSrcsByExtension(ctx, srcs)
+
+	for p, lla := range partitionedSrcs {
+		// if there are no sources, there is no need for headers
+		if lla.IsEmpty() {
+			continue
+		}
+		lla.Append(implementationHdrs)
+		partitionedSrcs[p] = lla
+	}
+
+	srcs = partitionedSrcs[cppSrcPartition]
+	cSrcs := partitionedSrcs[cSrcPartition]
+	asSrcs := partitionedSrcs[asSrcPartition]
+
 	absoluteIncludes.DeduplicateAxesFromBase()
 	localIncludes.DeduplicateAxesFromBase()
 
@@ -351,8 +371,6 @@
 		}
 	}
 
-	srcs, cSrcs, asSrcs := groupSrcsByExtension(ctx, srcs)
-
 	stlPropsByArch := module.GetArchVariantProperties(ctx, &StlProperties{})
 	for _, configToProps := range stlPropsByArch {
 		for _, props := range configToProps {