Move sharding functions for reuse

Move shardPaths and shardTests to android.ShardPaths and
android.ShardStrings for reuse in other packages.

Bug: 144948629
Test: m checkbuild
Change-Id: I868802872c73616b80f56cbf11f959c01a8b793a
Merged-In: I868802872c73616b80f56cbf11f959c01a8b793a
(cherry picked from commit 0a2f719bcaad397f27bcf52cd6604df6b5cec03b)

4 files changed, 134 insertions, 14 deletions

diff --git a/android/util.go b/android/util.go
index f9dce6fe..b17422df 100644
--- a/android/util.go
+++ b/android/util.go
@@ -243,3 +243,36 @@ func matchPattern(pat, str string) bool {
 	}
 	return strings.HasPrefix(str, pat[:i]) && strings.HasSuffix(str, pat[i+1:])
 }
+
+// ShardPaths takes a Paths, and returns a slice of Paths where each one has at most shardSize paths.
+func ShardPaths(paths Paths, shardSize int) []Paths {
+	if len(paths) == 0 {
+		return nil
+	}
+	ret := make([]Paths, 0, (len(paths)+shardSize-1)/shardSize)
+	for len(paths) > shardSize {
+		ret = append(ret, paths[0:shardSize])
+		paths = paths[shardSize:]
+	}
+	if len(paths) > 0 {
+		ret = append(ret, paths)
+	}
+	return ret
+}
+
+// ShardStrings takes a slice of strings, and returns a slice of slices of strings where each one has at most shardSize
+// elements.
+func ShardStrings(s []string, shardSize int) [][]string {
+	if len(s) == 0 {
+		return nil
+	}
+	ret := make([][]string, 0, (len(s)+shardSize-1)/shardSize)
+	for len(s) > shardSize {
+		ret = append(ret, s[0:shardSize])
+		s = s[shardSize:]
+	}
+	if len(s) > 0 {
+		ret = append(ret, s)
+	}
+	return ret
+}
diff --git a/android/util_test.go b/android/util_test.go
index 2e5eb07e..7f0d331c 100644
--- a/android/util_test.go
+++ b/android/util_test.go
@@ -404,3 +404,102 @@ func ExampleCopyOf_append() {
 	// b = ["foo" "bar"]
 	// c = ["foo" "baz"]
 }
+
+func Test_Shard(t *testing.T) {
+	type args struct {
+		strings   []string
+		shardSize int
+	}
+	tests := []struct {
+		name string
+		args args
+		want [][]string
+	}{
+		{
+			name: "empty",
+			args: args{
+				strings:   nil,
+				shardSize: 1,
+			},
+			want: [][]string(nil),
+		},
+		{
+			name: "single shard",
+			args: args{
+				strings:   []string{"a", "b"},
+				shardSize: 2,
+			},
+			want: [][]string{{"a", "b"}},
+		},
+		{
+			name: "single short shard",
+			args: args{
+				strings:   []string{"a", "b"},
+				shardSize: 3,
+			},
+			want: [][]string{{"a", "b"}},
+		},
+		{
+			name: "shard per input",
+			args: args{
+				strings:   []string{"a", "b", "c"},
+				shardSize: 1,
+			},
+			want: [][]string{{"a"}, {"b"}, {"c"}},
+		},
+		{
+			name: "balanced shards",
+			args: args{
+				strings:   []string{"a", "b", "c", "d"},
+				shardSize: 2,
+			},
+			want: [][]string{{"a", "b"}, {"c", "d"}},
+		},
+		{
+			name: "unbalanced shards",
+			args: args{
+				strings:   []string{"a", "b", "c"},
+				shardSize: 2,
+			},
+			want: [][]string{{"a", "b"}, {"c"}},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Run("strings", func(t *testing.T) {
+				if got := ShardStrings(tt.args.strings, tt.args.shardSize); !reflect.DeepEqual(got, tt.want) {
+					t.Errorf("ShardStrings(%v, %v) = %v, want %v",
+						tt.args.strings, tt.args.shardSize, got, tt.want)
+				}
+			})
+
+			t.Run("paths", func(t *testing.T) {
+				stringsToPaths := func(strings []string) Paths {
+					if strings == nil {
+						return nil
+					}
+					paths := make(Paths, len(strings))
+					for i, s := range strings {
+						paths[i] = PathForTesting(s)
+					}
+					return paths
+				}
+
+				paths := stringsToPaths(tt.args.strings)
+
+				var want []Paths
+				if sWant := tt.want; sWant != nil {
+					want = make([]Paths, len(sWant))
+					for i, w := range sWant {
+						want[i] = stringsToPaths(w)
+					}
+				}
+
+				if got := ShardPaths(paths, tt.args.shardSize); !reflect.DeepEqual(got, want) {
+					t.Errorf("ShardPaths(%v, %v) = %v, want %v",
+						paths, tt.args.shardSize, got, want)
+				}
+			})
+		})
+	}
+}
diff --git a/java/aapt2.go b/java/aapt2.go
index bcc8e976..f21408f9 100644
--- a/java/aapt2.go
+++ b/java/aapt2.go
@@ -61,7 +61,7 @@ var aapt2CompileRule = pctx.AndroidStaticRule("aapt2Compile",
 	"outDir", "cFlags")
 
 func aapt2Compile(ctx android.ModuleContext, dir android.Path, paths android.Paths) android.WritablePaths {
-	shards := shardPaths(paths, AAPT2_SHARD_SIZE)
+	shards := android.ShardPaths(paths, AAPT2_SHARD_SIZE)
 
 	ret := make(android.WritablePaths, 0, len(paths))
 
diff --git a/java/java.go b/java/java.go
index bf738c4d..9ac38c92 100644
--- a/java/java.go
+++ b/java/java.go
@@ -530,18 +530,6 @@ func hasSrcExt(srcs []string, ext string) bool {
 	return false
 }
 
-func shardPaths(paths android.Paths, shardSize int) []android.Paths {
-	ret := make([]android.Paths, 0, (len(paths)+shardSize-1)/shardSize)
-	for len(paths) > shardSize {
-		ret = append(ret, paths[0:shardSize])
-		paths = paths[shardSize:]
-	}
-	if len(paths) > 0 {
-		ret = append(ret, paths)
-	}
-	return ret
-}
-
 func (j *Module) hasSrcExt(ext string) bool {
 	return hasSrcExt(j.properties.Srcs, ext)
 }
@@ -1088,7 +1076,7 @@ func (j *Module) compile(ctx android.ModuleContext, extraSrcJars ...android.Path
 			shardSize := int(*(j.properties.Javac_shard_size))
 			var shardSrcs []android.Paths
 			if len(uniqueSrcFiles) > 0 {
-				shardSrcs = shardPaths(uniqueSrcFiles, shardSize)
+				shardSrcs = android.ShardPaths(uniqueSrcFiles, shardSize)
 				for idx, shardSrc := range shardSrcs {
 					classes := android.PathForModuleOut(ctx, "javac", jarName+strconv.Itoa(idx))
 					TransformJavaToClasses(ctx, classes, idx, shardSrc, nil, flags, extraJarDeps)