Refactor code for partitions c srcs

To support protos (and other srcs that generate sources), we need to
partition further. Separate out into a separate common function.

Bug: 200601772
Test: build/bazel/ci/bp2build.sh
Change-Id: I7bf4cd96fd9a9fca4ccb3c96f21a04303201f891
diff --git a/cc/bp2build.go b/cc/bp2build.go
index e48f757..c281c0e 100644
--- a/cc/bp2build.go
+++ b/cc/bp2build.go
@@ -14,12 +14,12 @@
 package cc
 
 import (
-	"fmt"
 	"path/filepath"
 	"strings"
 
 	"android/soong/android"
 	"android/soong/bazel"
+
 	"github.com/google/blueprint"
 
 	"github.com/google/blueprint/proptools"
@@ -41,18 +41,6 @@
 }
 
 func groupSrcsByExtension(ctx android.TopDownMutatorContext, srcs bazel.LabelListAttribute) (cppSrcs, cSrcs, asSrcs bazel.LabelListAttribute) {
-	// Branch srcs into three language-specific groups.
-	// 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.
-	// TODO(b/190006308): Handle language detection of sources in a Bazel rule.
-	isCSrcOrFilegroup := func(s string) bool {
-		return strings.HasSuffix(s, ".c") || strings.HasSuffix(s, "_c_srcs")
-	}
-
-	isAsmSrcOrFilegroup := func(s string) bool {
-		return strings.HasSuffix(s, ".S") || strings.HasSuffix(s, ".s") || strings.HasSuffix(s, "_as_srcs")
-	}
-
 	// Check that a module is a filegroup type named <label>.
 	isFilegroupNamed := func(m android.Module, fullLabel string) bool {
 		if ctx.OtherModuleType(m) != "filegroup" {
@@ -61,54 +49,39 @@
 		labelParts := strings.Split(fullLabel, ":")
 		if len(labelParts) > 2 {
 			// There should not be more than one colon in a label.
-			panic(fmt.Errorf("%s is not a valid Bazel label for a filegroup", fullLabel))
-		} else {
-			return m.Name() == labelParts[len(labelParts)-1]
+			ctx.ModuleErrorf("%s is not a valid Bazel label for a filegroup", fullLabel)
 		}
+		return m.Name() == labelParts[len(labelParts)-1]
 	}
 
-	// Convert the filegroup dependencies into the extension-specific filegroups
-	// filtered in the filegroup.bzl macro.
-	cppFilegroup := func(label string) string {
-		m, exists := ctx.ModuleFromName(label)
-		if exists {
-			aModule, _ := m.(android.Module)
-			if isFilegroupNamed(aModule, label) {
-				label = label + "_cpp_srcs"
+	// Convert filegroup dependencies into extension-specific filegroups filtered in the filegroup.bzl
+	// macro.
+	addSuffixForFilegroup := func(suffix string) bazel.LabelMapper {
+		return func(ctx bazel.OtherModuleContext, label string) (string, bool) {
+			m, exists := ctx.ModuleFromName(label)
+			if !exists {
+				return label, false
 			}
-		}
-		return label
-	}
-	cFilegroup := func(label string) string {
-		m, exists := ctx.ModuleFromName(label)
-		if exists {
 			aModule, _ := m.(android.Module)
-			if isFilegroupNamed(aModule, label) {
-				label = label + "_c_srcs"
+			if !isFilegroupNamed(aModule, label) {
+				return label, false
 			}
+			return label + suffix, true
 		}
-		return label
-	}
-	asFilegroup := func(label string) string {
-		m, exists := ctx.ModuleFromName(label)
-		if exists {
-			aModule, _ := m.(android.Module)
-			if isFilegroupNamed(aModule, label) {
-				label = label + "_as_srcs"
-			}
-		}
-		return label
 	}
 
-	cSrcs = bazel.MapLabelListAttribute(srcs, cFilegroup)
-	cSrcs = bazel.FilterLabelListAttribute(cSrcs, isCSrcOrFilegroup)
+	// 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")},
+		// 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},
+	})
 
-	asSrcs = bazel.MapLabelListAttribute(srcs, asFilegroup)
-	asSrcs = bazel.FilterLabelListAttribute(asSrcs, isAsmSrcOrFilegroup)
-
-	cppSrcs = bazel.MapLabelListAttribute(srcs, cppFilegroup)
-	cppSrcs = bazel.SubtractBazelLabelListAttribute(cppSrcs, cSrcs)
-	cppSrcs = bazel.SubtractBazelLabelListAttribute(cppSrcs, asSrcs)
+	cSrcs = partitioned["c"]
+	asSrcs = partitioned["as"]
+	cppSrcs = partitioned["cpp"]
 	return
 }